from django.core.management.base import BaseCommand, CommandError
from monthplan.models import SurveyHealpixPlate, NSIDE_PLATES, ORDER, MAPS_DIR, Survey, SurveyPath

from os import listdir
from astropy_healpix import HEALPix
from astropy import units as u
from astropy.coordinates import SkyCoord, RangeError
from astropy.coordinates import ICRS, Galactic, FK5
from astropy.coordinates import Angle

import numpy as np

from srglib.utils import slerp, quat_to_pol_and_roll
import math
from Quaternion import Quat

class Command(BaseCommand):

    
    def handle(self, *args, **options):
        tolerance = 10.0
        hp = HEALPix(nside=NSIDE_PLATES, order=ORDER, frame=FK5())
        surveys=Survey.objects.all()
        for survey in surveys:
            print(survey.experiment)
            surveypaths = survey.surveypath_set.all()
            for i in range(len(surveypaths)-1):
                delta=(surveypaths[i+1].dtime-surveypaths[i].dtime).total_seconds()
                if(abs(int(delta)) > tolerance):
                    nstep = int(abs(delta) / tolerance)
                    delta_arr = np.linspace(0,1,nstep)

                    quat = slerp([surveypaths[i].q2, surveypaths[i].q3, surveypaths[i].q4, surveypaths[i].q1],
                                 [surveypaths[i+1].q2, surveypaths[i+1].q3, surveypaths[i+1].q4, surveypaths[i+1].q1],
                                 delta_arr)
                    for q in quat:
                        qfin=Quat(attitude=q)
                        crd = SkyCoord(qfin.ra, qfin.dec, frame="fk5", unit="deg")                                                                                                                            
                        heal = hp.skycoord_to_healpix(crd)
                        try:
                            plate = SurveyHealpixPlate.objects.get(healpix=heal)
                        except:
                            print('Error: SurveyHealpixPlate not found, run ./manage.py init_survey_healpix_plates first')
                            break
                        queryset = survey.surveyhealpixplate_set.all()
                        if not queryset.filter(pk=plate.pk).exists():
                            plate.survey.add(survey)
                            plate.save()