# -*- encoding: utf-8 -*-
'''
@File    :   scampy.py
@Time    :   2022/09/10
@Author  :   Mingyu Li
@Contact :   lmytime@hotmail.com
'''


from astropy.io import fits
from astropy.wcs import WCS
from astrolib.astro_image import AstroImage
import getpass
import subprocess
import os, time

def dict_to_command_string(config):
    return "".join([ f" -{key} {value} " for key, value in config.items() ])

class Sextrapy:
    def __init__(self, file, folder="."):
        '''
            Initialize the class, specify the folder to store the files
        '''
        self.dir = folder
        self.file = file
        os.makedirs(self.dir, exist_ok=True)

    @classmethod
    def config_sextractor_for_scamp(self):
        sextractor_config = {
            "CATALOG_TYPE": "FITS_LDAC",
            "DETECT_THRESH": 3.0,
            "ANALYSIS_THRESH": 3.0,
            "FILTER": "N",
            "PIXEL_SCALE": 0.0, # size of pixel in arcsec (0=use FITS WCS info)
            "MAG_ZEROPOINT": 0.0,
        }
        return sextractor_config

    @classmethod
    def config_sextractor(self):
        sextractor_config = {
            "CATALOG_NAME": "test.cat",
            "CATALOG_TYPE": "ASCII_HEAD",
            "PARAMETERS_NAME": "default.param",
            "DETECT_TYPE": "CCD",
            "DETECT_MINAREA": 5,
            "DETECT_THRESH": 3,
            "ANALYSIS_THRESH": 3,
            "FILTER": "Y",
            "FILTER_NAME": "default.conv",
            "DEBLEND_NTHRESH": 32,
            "DEBLEND_MINCONT": 0.005,
            "CLEAN": "Y",
            "CLEAN_PARAM": 1.0,
            "MASK_TYPE": "CORRECT",
            "PHOT_APERTURES": 5,
            "PHOT_AUTOPARAMS": "2.5,3.5",
            "PHOT_PETROPARAMS": "2.0,3.5",
            "SATUR_LEVEL": 50000.0,
            "SATUR_KEY": "SATURATE",
            "MAG_ZEROPOINT": 0.0,
            "MAG_GAMMA": 4.0,
            "GAIN": 0.0,
            "GAIN_KEY": "GAIN",
            "PIXEL_SCALE": 0.0,
            "SEEING_FWHM": 1.2,
            "STARNNW_NAME": "default.nnw",
            "BACK_SIZE": 64,
            "BACK_FILTERSIZE": 3,
            "BACKPHOTO_TYPE": "GLOBAL",
            "CHECKIMAGE_TYPE": "NONE",
            "CHECKIMAGE_NAME": "check.fits",
            "MEMORY_OBJSTACK": 3000,
            "MEMORY_PIXSTACK": 300000,
            "MEMORY_BUFSIZE": 1024,
            "VERBOSE_TYPE": "NORMAL",
            "HEADER_SUFFIX": ".head",
            "WRITE_XML": "N",
        }
        return sextractor_config

    def paramters_sextractor_for_scamp(self):
        '''
            Write a paramters_sextractor_for_scamp.param file for sextractor
        '''
        params = ["NUMBER", "FLUX_AUTO", "FLUXERR_AUTO", "ALPHAWIN_J2000", "DELTAWIN_J2000",
                    "ERRAWIN_WORLD", "ERRBWIN_WORLD", "ERRTHETA_WORLD",
                    "XWIN_IMAGE", "YWIN_IMAGE", "ERRAWIN_IMAGE", "ERRBWIN_IMAGE",
                    "ERRTHETAWIN_IMAGE", "MAG_AUTO", "MAGERR_AUTO", "FLAGS", "FLAGS_WEIGHT"]
        params = "\n".join(params)
        params = "#This file is generated by astrolib.astromatic module\n#      (c) Mingyu\n# Parameters prepared for SCAMP\n" + params
        with open(os.path.join(self.dir, "paramters_sextractor_for_scamp.param"), "w") as f:
            f.write(params)

    def run_sextractor(self, sextractor="sex", extension=None, config=None):
        # sepcify sextractor configure
        if config is None:
            sextractor_config = self.config_sextractor()
        else:
            sextractor_config = config

        # specify extension of fits file to be processed
        if extension is None:
            extension_string = ""
        else:
            extension_string = f"[{int(extension)}]"

        # run sextractor
        config_string = dict_to_command_string(sextractor_config)
        command_string = f"{sextractor} {self.file}{extension_string} {config_string}"
        print(command_string)
        subprocess.call(command_string, shell=True)


class Scampy:

    def __init__(self, file, folder="."):
        '''
            Initialize the class, specify the folder to store the files
        '''
        self.dir = folder
        self.file = file
        os.makedirs(self.dir, exist_ok=True)

    def run_sextractor_for_scamp(self, sextractor="sex", extension=None, config=None):
        sex = Sextrapy(self.file, self.dir)
        # sepcify sextractor configure
        if config is None:
            config = sex.config_sextractor_for_scamp()
        config['CATALOG_NAME'] = os.path.join(self.dir, self.file.replace('.fits', '.LDAC'))
        config['PARAMETERS_NAME'] = os.path.join(self.dir, 'paramters_sextractor_for_scamp.param')
        sex.paramters_sextractor_for_scamp()
        sex.run_sextractor(sextractor=sextractor, extension=None, config=config)

    @classmethod
    def config_scamp(self):
        scamp_config = {
            'REF_SERVER': 'vizier.nao.ac.jp', # Internet addresses of catalog servers
                                                # Possible mirrors include:
                                                # vizier.nao.ac.jp,
                                                # vizier.hia.nrc.ca,
                                                # vizier.ast.cam.ac.uk,
                                                # vizier.iucaa.in,
                                                # vizier.china-vo.org,
                                                # vizier.cfa.harvard.edu and
                                                # viziersaao.chpc.ac.za
            'ASTREF_CATALOG': 'GAIA-DR2',   # NONE,FILE,USNO-A2,USNO-B1,GSC-2.3,
                                                # TYCHO-2,UCAC-4,URAT-1,NOMAD-1,PPMX,
                                                # CMC-15,2MASS,DENIS-3,SDSS-R9,SDSS-R12,
                                                # IGSL,GAIA-DR1,GAIA-DR2,GAIA-EDR3,
                                                # PANSTARRS-1, or ALLWISE
            'ASTREF_BAND': 'DEFAULT',
            'ASTREFCAT_NAME': 'astrefcat.cat',
            'ASTREFCENT_KEYS': 'ALPHAWIN_J2000,DELTAWIN_J2000',
            'ASTREFERR_KEYS': 'ERRAWIN_WORLD,ERRBWIN_WORLD,ERRTHETAWIN_WORLD',
            'ASTREFMAG_KEY': 'MAG_AUTO',
            'ASTREFMAGERR_KEY': 'MAGERR_AUTO',
            'ASTREFMAG_LIMITS': '16.0,19.0', # Select magnitude range in ASTREF_BAND
            'CROSSID_RADIUS': '1.0',
            'DISTORT_DEGREES': '2',
            'SOLVE_PHOTOM': 'N',
            'SN_THRESHOLDS': '20.0,100.0',   # S/N thresholds (in sigmas) for all and
                                                # high-SN sample
            'AHEADER_SUFFIX': '.ahead',
            'HEADER_SUFFIX': '.head',
            'WRITE_XML': 'N',
        }
        return scamp_config

    def run_scamp(self, scamp="scamp", config=None):
        # sepcify sextractor configure
        if config is None:
            scamp_config = self.config_scamp()
        else:
            scamp_config = config

        # run scamp
        config_string = dict_to_command_string(scamp_config)
        command_string = f"{scamp} {os.path.join(self.dir, self.file.replace('.fits', '.LDAC'))} {config_string}"
        print(command_string)
        subprocess.call(command_string, shell=True)

    @classmethod
    def read_scamp_header(self, header=None):
        '''
            Read the header file generated by scamp
        '''
        if header is None:
            header = os.path.join(self.dir, self.file.replace('.fits', '.head'))
        with open(header, "r") as f:
            headers = f.readlines()
        new_header = fits.Header.fromstring("".join(headers[3:17]))
        return new_header

    @classmethod
    def update_header(self, header, new_header):
        '''
            Update the header of the original header
        '''

        for card in new_header.cards:
            header[card[0]] = card[1:]
        header['history'] = f"WCS calibrated with SCAMP by {getpass.getuser()} on {time.strftime('%Y/%m/%d %H:%m', time.localtime())}"

        return header

    @classmethod
    def update_header_for_fits(self, extension, new_header, file=None):
        '''
            Update the header of the original fits file
        '''
        if file is None:
            file = self.file
        with fits.open(file) as hdul:
            header = hdul[extension].header
            hdul[extension].header = self.update_header(header, new_header)
            hdul.writeto(file.replace('.fits', '_newwcs.fits'), overwrite=True)