from django.core.management.base import BaseCommand, CommandError

import os
import io
import gzip
import tarfile
import warnings

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

from avro.datafile import DataFileReader, DataFileWriter
from avro.io import DatumReader, DatumWriter
import fastavro

from astropy.time import Time
from astropy.time import Time, TimezoneInfo
from astropy.io import fits
import astropy.units as u
import aplpy

from astropy.time import Time
from django.utils.timezone import localdate, localtime, now

from collections import defaultdict

import logging
import requests 
import json
import pprint

from astrobasis.models import ALeRCE
from srglib.ztf import find_ztf_in_survey, load_ztf_alerce

from astropy.coordinates import SkyCoord # High-level coordinates
from astropy.coordinates import ICRS, Galactic, FK4, FK5 # Low-level frames
from astropy.coordinates import Angle, Latitude, Longitude # Angles
import astropy.units as u
import re
import json
import pandas as pd
from astropy_healpix import HEALPix
from erotrans.models import EroTransSource, EroTrans
from erosurvey.models import NSIDE_PLATES, ORDER_PLATES
from heasarc.models import NSIDE_SOURCES, ORDER
from monthplan.models import Survey

def load_ztf_alerce_to_django_local(item):
    utc_moscow = TimezoneInfo(utc_offset=3*u.hour)
    hp = HEALPix(nside=NSIDE_SOURCES,
                 order=ORDER, frame=FK5())

    hp_plates = HEALPix(nside=NSIDE_PLATES,
                 order=ORDER_PLATES, frame=FK5())

    try:
        ztf = ALeRCE.objects.get(oid=item['oid'])
        print("ALeRCE ZTF alert %s is already loaded, skip" % item['oid'])
        return
    except ALeRCE.DoesNotExist:
        print("ALeRCE ZTF alert %s is not loaded" % item['oid'])
        pass
    
    ra = float(item['meanra'])
    dec = float(item['meandec'])
    crd = SkyCoord(ra, dec, frame=FK5(), unit="deg")
    healpix = hp.skycoord_to_healpix(crd)
    healpix_plate = hp_plates.skycoord_to_healpix(crd)

    first_tm = Time(item['firstmjd'], format='mjd', scale='utc')
    last_tm  = Time(item['lastmjd'], format='mjd', scale='utc')
    first_dt = first_tm.to_datetime(timezone=utc_moscow)
    last_dt  = last_tm.to_datetime(timezone=utc_moscow)

    ztf = ALeRCE(
        healpix  = healpix,
        healpix_plate  = healpix_plate,
        classearly = item['classearly'],
        classrf = item['classrf'],
        oid = item['oid'],
        firstdate = first_dt,
        lastdate = last_dt,
        firstmjd = item['firstmjd'],
        lastmjd = item['lastmjd'],
        mean_magap_g = item['mean_magap_g'],
        mean_magap_r = item['mean_magap_r'],
        mean_magpsf_g = item['mean_magpsf_g'],
        mean_magpsf_r = item['mean_magpsf_r'],
        dec = dec,
        ra = ra,
        nobs = item['nobs'],
        pclassearly = item['pclassearly'],
        pclassrf = item['pclassrf'],
        sigma_magap_g = item['sigma_magap_g'],
        sigma_magap_r = item['sigma_magap_r'],
        sigma_magpsf_g = item['sigma_magpsf_g'],
        sigma_magpsf_r = item['sigma_magpsf_r'],
        sigmadec = item['sigmadec'],
        sigmara = item['sigmara'])
        
    ztf.save()
    find_ztf_in_survey(ztf)

def load_ztf_alerce_local(days, extra=None):
    classrf = "sn ia"
    pclassrf = 0.1
    pclassearly = 0.1
    sortBy = "firstmjd"
    nobsmin = 2
    nobsmax = 40
    url = "https://ztf.alerce.online/query"
    
    current_utc = now()
    current_mjd = Time(current_utc).mjd
    firstmjd_min = current_mjd - days

    post={
        "sortBy": sortBy,
        "records_per_pages": 100,
        "query_parameters":{
            "filters": {
            "nobs": { 
                "min": nobsmin,
                "max": nobsmax
            },
            #"classrf": classrf,
            #"pclassrf": pclassrf,
            'pclassearly':pclassearly,
            },
            "dates":{
                "firstmjd":{
                    "min": firstmjd_min
                }
            }
        }
    }
    
    if(extra):
        post['query_parameters']['filters'].update(extra)

    #pprint.pprint(post)
    
    r = requests.post(url = url, json = post) 
    data = r.json() 
    for item in data['result']:
        """ loads first page """
        packet=data['result'][item]
        #pprint.pprint(packet['oid'])
        load_ztf_alerce_to_django_local(packet)

    total=int(data['total'])
    num_pages=int(data['num_pages'])
    page=int(data['page'])
    
    print('total ',total,'num_pages ',num_pages,'page ', page)
    #return
    pages = list(range(2,num_pages+1))
    for page in pages:
        """ loads all other pages, staring from 2 """
        post.update( {'page':page,} )
        pprint.pprint(post)
        r = requests.post(url = url, json = post) 
        data = r.json() 
        for item in data['result']:
            packet=data['result'][item]
            #pprint.pprint(packet['oid'])
            load_ztf_alerce_to_django_local(packet)

class Command(BaseCommand):
    help = 'Initiates data dase'

    def handle(self, *args, **options):
        #all = ALeRCE.objects.all()
        #all.delete()

        load_ztf_alerce_local(2)

        print('Done')