import warnings
warnings.filterwarnings('ignore')
# import sys
# import os

# import numpy as np
import ccdproc
from astropy.time import Time
from astropy.table import QTable, Table
import astroalign as aa
from astropy.wcs import WCS
from astropy.utils.console import ProgressBar, ProgressBarOrSpinner
# from astropy.coordinates import SkyCoord as acoord
# import astropy.units as un
from astropy.io import fits

from astropy.nddata.ccddata import CCDData

# photometry imports
from photutils.psf import IntegratedGaussianPRF, DAOGroup
from photutils.background import MMMBackground, MADStdBackgroundRMS
from astropy.modeling.fitting import LevMarLSQFitter
from astropy.stats import gaussian_sigma_to_fwhm
from photutils.psf import IterativelySubtractedPSFPhotometry
from photutils import DAOStarFinder
from astropy.stats import mad_std
# from astroquery.simbad import Simbad

'''
Ceres is the handler of series in Dorado,
'''

# __all__ = [Ceres]

class Ceres:

    def __init__(self, filters = {}, data = [], bias = None, time = None, datestr = None):
        # metadata
        self.filters = filters
        self.data = data
        self.bias = bias
        self.time = time
        self.datestr = datestr
        
        # location, weather, timezone, camera, observer array

        # refering instance of clippy to call and save later?
        # or call clippy directly and feed it a ceres object

        try:
            self.date = Time(int(self.time.mjd), format = 'mjd')
        except:
            self.date = None
        
    def add_stack(self, stack):
        # eventually stacks themelves should have some metadata 
        # to denote stuff like calibration status
        self.filters[stack.filter] = len(self.data)
        self.data.append(stack)

    def rem_stack(self, filter):
        del self.data[self.filters[filter]]
        # delete time strings

    def calibrate(self, filter):
        # for bla in series: add bias corrected = True to header
        stack = self.data[self.filters[filter]]
        flat = stack.flat
        bias = self.bias
        c_series = []
        with ProgressBar(len(stack.data)) as bar:
            for im in stack.data:
                bar.update()
                im.data = im.data.astype('uint16') 
                flat.data = flat.data.astype('uint16') 
                bias.data = bias.data.astype('uint16') 
                im = ccdproc.ccd_process(im, master_bias = bias, master_flat = flat)
                im.data = im.data.astype('uint16') 
                c_series.append(im)
        self.data[self.filters[filter]].data = c_series
        self.data[self.filters[filter]].calibrated = True

    def imarith(self, filter, operator, operand):
        # mod to check datatype using type()
        # mod to remove im_count and make possible to use single image
        # mod to accomodate CCDdata object
        series = self.data[self.filters[filter]]
        for i in range(len(series)):
            if (operator == '+'):
                series[i].data = series[i].data  + operand
            elif (operator == '-'):
                series[i].data = series[i].data - operand
            elif (operator == '/'):
                series[i].data = series[i].data  / operand
            elif (operator == '*'):
                series[i].data = series[i].data  * operand
        
        self.data[self.filters[filter]] = series

    def getWCS(self, filter, clippy, alignto = None, cache = True):
        series = self.data[self.filters[filter]]
        if alignto == None:
            alignto = series.alignTo
        if cache:
            hdulist = fits.open(clippy.dordir / 'cache' / 'astrometryNet' / 'solved.fits') 
            self.data[self.filters[filter]].wcs = WCS(hdulist[0].header, hdulist)
            self.data[self.filters[filter]].solved = CCDData.read(clippy.dordir / 'cache' / 'astrometryNet' / 'solved.fits')
            hdulist.close()
        else:
            toalign = series.data[alignto]
            fname, cachedir = clippy.mkcacheObj(toalign, 'astrometryNet')
            path = [cachedir, fname]
            writearray = [cachedir, 'solved.fits']
            solved, wcs_header = clippy.plate_solve(path, writearray = writearray)
            clippy.delcacheObj( fname, 'astrometryNet')
            self.data[self.filters[filter]].wcs = WCS(wcs_header)
            self.data[self.filters[filter]].solved = solved

    def align(self, filter, clippy, alignto = None, getWCS = True, cache = False):
        series = self.data[self.filters[filter]]
        if alignto == None:
            alignto = series.alignTo
        toalign = series.data[alignto]
        if getWCS:
            if cache:
                toalign =  CCDData.read(clippy.dordir / 'cache' / 'astrometryNet' / 'solved.fits', unit = clippy.unit)
                hdulist = fits.open(clippy.dordir / 'cache' / 'astrometryNet' / 'solved.fits') 
                self.data[self.filters[filter]].wcs = WCS(hdulist[0].header, hdulist)
                hdulist.close()
                self.data[self.filters[filter]].solved = toalign
            else:
                fname, cachedir = clippy.mkcacheObj(toalign, 'astrometryNet')
                path = [cachedir, fname]
                writearray = [cachedir, 'solved.fits']
                solved, wcs_header = clippy.plate_solve(path, writearray = writearray)
                toalign = solved
                clippy.delcacheObj( fname, 'astrometryNet')
                self.data[self.filters[filter]].wcs = WCS(wcs_header)
                self.data[self.filters[filter]].solved = solved
                # delete cache object
                # save solved to target

        aa_series = []
        with ProgressBar(len(series.data)) as bar:
            for image in series.data:
                bar.update()
                try:
                    img, _ = aa.register(image.data, toalign.data)
                    image.data = img
                    aa_series.append(image)
                except:
                    print('Image skipped')

        self.data[self.filters[filter]].data = aa_series
        self.data[self.filters[filter]].aligned = True

    def dorphot(self, filter, zellars):
        # get seeing from PSF
        stack = self.data[self.filters[filter]]
        # if no wcs, complain alot
        w = stack.wcs
        xy = w.wcs_world2pix(zellars.coords.ra.deg, zellars.coords.dec.deg, 1)

        print(float(xy[0]), ' : ',  float(xy[1]))

        pos = Table(names=['x_0', 'y_0'], data=xy)
        sigma_psf = 2.0
        bkg_sigma = mad_std(stack.data[0]) 
        daofind = DAOStarFinder(fwhm = 4., threshold = 3. * bkg_sigma)  
        daogroup = DAOGroup(2.0 * sigma_psf * gaussian_sigma_to_fwhm)
        mmm_bkg = MMMBackground()

        psf_model = IntegratedGaussianPRF(sigma = sigma_psf)
        psf_model.x_0.fixed = True
        psf_model.y_0.fixed = True

        fitter = LevMarLSQFitter()
        bkgrms = MADStdBackgroundRMS()

        times = []
        exptimes = []
        ray = []
        decx = []
        x = []
        y = []
        flux = []
        fluxunc = []
        # if radec get xy
        with ProgressBar(len(stack.data)) as bar:
            for image in stack.data:
                bar.update()
                photometry = IterativelySubtractedPSFPhotometry(finder = daofind, group_maker = daogroup, bkg_estimator = mmm_bkg,
                        psf_model = psf_model, fitter = LevMarLSQFitter(), niters = 1, fitshape = (21, 21))
                results = photometry(image = image, init_guesses= pos)
                [ra, dec] = w.wcs_pix2world(results['x_fit'], results['y_fit'], 1)
                
                times.append(Time(image.header['DATE-OBS']))
                exptimes.append(image.header['EXPTIME'])
                ray.append(ra)
                decx.append(dec)
                x.append(results['x_fit'])
                y.append(results['y_fit'])
                flux.append(results['flux_fit'])
                fluxunc.append(results['flux_unc'])
                # if ts:
                    # ts.add_row([time, exptime, results['x_fit'], results['y_fit'], ra, dec, results['flux_fit'], results['flux_unc']])
                # else:
                    # ts = QTable([time, exptime, results['x_fit'], results['y_fit'], ra, dec, results['flux_fit'], results['flux_unc']], names=('time', 'exptime', 'x', 'y', 'ra', 'dec', 'flux', 'flux_unc'), meta={'name': filter})
                # instrumental magnitude
        ts = QTable([times, exptimes, x, y, ray, decx, flux, fluxunc], names=('time', 'exptime', 'x', 'y', 'ra', 'dec', 'flux', 'flux_unc'), meta={'name': filter})
        
        zellars.filters[filter] = len(zellars.ts)
        zellars.ts.append(ts)

    def dorphotc(self, filter, zellars, zcontrol, shape = 21):
        # get seeing from PSF
        stack = self.data[self.filters[filter]]
        # if no wcs, complain alot
        w = stack.wcs
        xy = w.wcs_world2pix(zellars.coords.ra.deg, zellars.coords.dec.deg, 1)
        xyc = w.wcs_world2pix(zcontrol.coords.ra.deg, zcontrol.coords.dec.deg, 1)
        
        # print(float(xy[0]), ' : ',  float(xy[1]))
        # print(xyc)

        pos = Table(names=['x_0', 'y_0'], data = ([float(xy[0]), float(xyc[0])], [float(xy[1]), float(xyc[1])]))
        # print(pos)
        sigma_psf = 2.0
        bkg_sigma = mad_std(stack.data[0]) 
        daofind = DAOStarFinder(fwhm = 4., threshold = 3. * bkg_sigma)  
        daogroup = DAOGroup(2.0 * sigma_psf * gaussian_sigma_to_fwhm)
        mmm_bkg = MMMBackground()

        psf_model = IntegratedGaussianPRF(sigma = sigma_psf)
        psf_model.x_0.fixed = True
        psf_model.y_0.fixed = True

        fitter = LevMarLSQFitter()
        bkgrms = MADStdBackgroundRMS()

        times = []
        exptimes = []
        ray = []
        decx = []
        x = []
        y = []
        flux = []
        fluxunc = []
        # if radec get xy
        itera = 0
        with ProgressBarOrSpinner(len(stack.data), msg = 'Performing PSF photometry') as bar:
            for image in stack.data:
                bar.update(itera)
                itera = itera + 1
                photometry = IterativelySubtractedPSFPhotometry(finder = daofind, group_maker = daogroup, bkg_estimator = mmm_bkg,
                        psf_model = psf_model, fitter = LevMarLSQFitter(), niters = 1, fitshape = (shape, shape))
                results = photometry(image = image, init_guesses= pos)
                [ra, dec] = w.wcs_pix2world(results['x_fit'], results['y_fit'], 1)
                
                times.append(Time(image.header['DATE-OBS']))
                exptimes.append(image.header['EXPTIME'])
                ray.append(ra)
                decx.append(dec)
                x.append(results['x_fit'][0])
                y.append(results['y_fit'][0])
                flux.append(results['flux_fit'][0] - results['flux_fit'][1])
                fluxunc.append(results['flux_unc'][0])

        ts = QTable([times, exptimes, x, y, ray, decx, flux, fluxunc], names=('time', 'exptime', 'x', 'y', 'ra', 'dec', 'flux', 'flux_unc'), meta={'name': filter})
        
        zellars.filters[filter] = len(zellars.ts)
        zellars.ts.append(ts)


# class Stack:
#     def __init__(self, data, flat = None, filter = '', times = [], calibrated = None, aligned = None, target = None, alignTo = 0):
#         self.data = data
#         self.flat = flat
#         self.filter = filter
#         self.length = len(data)
#         self.calibrated = calibrated
#         self.aligned = aligned

#         self.target = target
#         self.target_info = {} # dictionary of values, call simbad?

#         self.times = times 
#         self.wcs = None
#         self.alignTo = alignTo
#         self.solved = None
#         # include things like flux uncertainty etc.
#         # save data
#         # if flat is none, call clippy for help

#         if self.filter == '':
#             try:
#                 self.filter = data[0].header['filter']
#             except:
#                 self.filter = ''
        
#         if self.times == []:
#             try: 
#                 self.get_times()
#                 print('Times set')
#             except:
#                 self.times = []
     
#     def get_times(self):
#         times = []
#         for im in self.data:
#             times.append(Time(im.header['DATE-OBS'], format='fits'))
#         self.times = times

#     def get_target_info(self, target = None):
#         if target != None:
#             self.target = target
        

# class zellars:
#     def __init__(self, name):
#         # get it because zellars is the canadian target?
#         self.name = name
#         s = Simbad()
#         r = s.query_object(self.name)
#         self.filters = {}
#         self.ts = []
#         # r.pprint()
#         # print(r.colnames)
#         self.coords = acoord(ra = r['RA'], dec = r['DEC'], unit = (un.hourangle, un.deg), frame = 'icrs')
#         # xy coordinates
        
#     def calcmag(self, filter):
#         flux  = self.ts[self.filters[filter]]['flux']
#         flux_unc = self.ts[self.filters[filter]]['flux_unc']
#         magnitudes = -2.5 * np.log10(flux / 15)
#         mag_unc = flux_unc / ((flux / 15) * 2.30258509)
#         self.ts[self.filters[filter]]['mag'] = magnitudes
#         self.ts[self.filters[filter]]['mag_unc'] = mag_unc

#     def record(self, clippy, cr, saveType = 'fits'):
#         wrkdir = clippy.dordir / 'data' / 'wrk'
#         if cr.datestr == None:
#             cr.datestr = clippy.getDateString(cr)
#         datestr = cr.datestr
#         wrdir = wrkdir / datestr
#         clippy.mkwrk(cr)
#         for fi in self.filters.keys():
#             wrts = self.ts[self.filters[fi]]
#             fname = str(self.name) + '_' + str(fi) + '-' + str(int(cr.date.mjd)) + '.' + saveType
#             wrts.write(wrdir / fname, overwrite = True)
        
#     def export(self, clippy, cr, objectClass = None):
#         tardir = clippy.dordir / 'targets'
#         if objectClass:
#             os.makedirs(tardir / objectClass, exist_ok = True)
#             tardir = tardir / objectClass
        
