#!/usr/bin/env python3

from argparse import ArgumentParser
import numpy as np
import matplotlib.pyplot as pl
from allium import Sphere

parser = ArgumentParser()
parser.add_argument('-c', type=float, nargs='+', default=[1.0])
parser.add_argument('-R', type=float, nargs='+', default=[1.0])
parser.add_argument('-l', type=int, nargs='+', default=[0])
args = parser.parse_args()

s = Sphere(c=args.c, R=args.R)

fig, axs = pl.subplots(2, 1, gridspec_kw={'height_ratios': [3,1]})
axs[1].step([0]+s.R, [s.c[0]]+s.c, 'k-')
axs[1].set_xlabel('r')
axs[1].set_ylabel('c')

f = np.arange(0.001, 20.*np.pi, 0.25*np.pi)
Delta = s.Delta_omega
s.search_eigenfrequencies(args.l, f)
    
for l in args.l:
    o_D = np.array(s[l])/Delta
    axs[0].plot(np.mod(o_D-0.01,1.0)+0.01, o_D, 'C%io' % l, label='ℓ=%i' % l)

axs[0].set_xlabel('mod(ω/Δ, 1), Δ=%.3fπ' % (Delta/np.pi))
axs[0].set_ylabel('ω/Δ')
axs[0].legend()
fig.tight_layout()

pl.show()
