import matplotlib.pyplot as plt
from astropy.time import Time
from mskpy import KeplerState
import cometsuite as cs
import cometsuite.generators as gen
import cometsuite.scalers as sc
date = Time("2011-09-11")
rc = [1.99512556e8, -1.86957354e8, 1.49657485e8]
vc = [-23.28613574, 10.86086487, 13.85289671]
re = [1.47206597e8, -3.19139879e7, 1.38953505e2]
ve = [5.82028904e+00,  2.89897439e+01, -1.16199526e-03]
comet = KeplerState(rc, vc, date, name="C/2009 P1")
earth = KeplerState(re, ve, date, name="Earth")
pgen = cs.Coma(comet, date)
pgen.composition = cs.Geometric(rho0=1)
pgen.age = gen.Uniform(0, 5)
pgen.radius = gen.Log(-1, 3)
pgen.vhat = gen.Isotropic()
pgen.speed = gen.Delta(0.3)
pgen.speed_scale = sc.SpeedRh(-0.5) * sc.SpeedRadius(-0.5)
pgen.nparticles = 2000
integrator = cs.BulirschStoer()
sim = cs.run(pgen, integrator)
sim.observer = earth
sim.observe()
camera = cs.Camera(size=(60, 60), cdelt=[-1, 1])
camera.integrate(sim)
fig, ax = plt.subplots(figsize=(6, 6), dpi=200)
ax.imshow(camera.data, vmin=0, vmax=50, cmap="gray_r", extent=[30, -30, -30, 30], origin="lower")
plt.setp(ax, xlabel="RA offset (arcsec)", ylabel="Dec offset (arcsec)")
plt.tight_layout()