
#%%

import subprocess
from datetime import datetime
import os
import glob
from typing import List, Optional, Union
from pathlib import Path
from ezphot.utils.tract7dt.formatter import Formatter
from ezphot.utils.tract7dt.configuration import Configuration
from ezphot.utils.tiles import Tiles
#%%

class Tract7DTRunner:
    

    
    def __init__(self,
                 image_paths: List[Path],
                 filter_list: List[str],
                 catalog_paths: Optional[List[Path]] = None):
        self.image_paths = image_paths
        self.filter_list = filter_list
        self.catalog_paths = catalog_paths
        self.formatter = Formatter(image_paths = image_paths,
                                   filter_list = filter_list,
                                   catalog_paths = catalog_paths)
        self.configuration = Configuration()

        self.id = datetime.now().strftime("%Y%m%d_%H%M")
        print(f"ID: {self.id}")
        self.workdir = Path.home() / 'tract7dt' / self.id
        print('Tract7DT working directory: ', self.workdir)
        self._register_configuration_paths()
        
    def __repr__(self):
        return f"Tract7DTRunner(id = {self.id}, workdir = {self.workdir}, image_list_saved = {self.input_images_saved}, catalog_list_saved = {self.input_catalogs_saved})"
    
    
    def _register_configuration_paths(self):
        self.configuration.inputs.input_catalog = self.input_catalogs_path
        self.configuration.inputs.image_list_file = self.input_images_path
        self.configuration.outputs.work_dir = self.workdir    
        
    @property
    def input_images_saved(self):
        return self.input_images_path.exists()
    
    @property
    def input_catalogs_saved(self):
        return self.input_catalogs_path.exists()
    
    @property
    def reference_catalog_saved(self):
        return Path(self.configuration.inputs.gaiaxp_synphot_csv).exists()
    
    @property
    def input_images_path(self):
        return self.workdir / f"input_images.txt"

    @property
    def input_catalogs_path(self):
        return self.workdir / f"input_catalogs.csv"
    
    @property
    def configuration_path(self):
        return self.workdir / f"configuration.yaml"
        
    def register_target(self, 
                        list_ra: List[float],
                        list_dec: List[float],
                        list_type: List[str] = None,
                        list_ellip: List[float] = None,
                        list_Re: List[float] = None,
                        list_theta: List[float] = None,
                        
                        update_type_from_catalog: bool = True,
                        update_flux_from_catalog: bool = True,
                        update_ellip_from_catalog: bool = True,
                        update_Re_from_catalog: bool = True,
                        update_theta_from_catalog: bool = True,
                        
                        ra_key: str = 'X_WORLD',
                        dec_key: str = 'Y_WORLD',
                        type_key: str = 'CLASS_STAR',
                        ellip_key: str = 'ELLIPTICITY',
                        flux_key: str = 'FLUX_AUTO',
                        Re_key: str = 'FLUX_RADIUS',
                        theta_key: str = 'THETA_IMAGE',
                        matching_radius_arcsec: float = 5):
        
        self.input_images_path.parent.mkdir(parents=True, exist_ok=True)
        self.input_catalogs_path.parent.mkdir(parents=True, exist_ok=True)
        
        self.formatter.to_input_images(output_path = self.input_images_path)
        
        self.formatter.to_input_catalogs(output_path = self.input_catalogs_path,
                                         list_ra = list_ra,
                                         list_dec = list_dec,
                                         list_type = list_type,
                                         list_ellip = list_ellip,
                                         list_Re = list_Re,
                                         list_theta = list_theta,
                                         update_type_from_catalog = update_type_from_catalog,
                                         update_flux_from_catalog = update_flux_from_catalog,
                                         update_ellip_from_catalog = update_ellip_from_catalog,
                                         update_Re_from_catalog = update_Re_from_catalog,
                                         update_theta_from_catalog = update_theta_from_catalog,
                                         ra_key = ra_key,
                                         dec_key = dec_key,
                                         type_key = type_key,
                                         ellip_key = ellip_key,
                                         flux_key = flux_key,
                                         Re_key = Re_key,
                                         theta_key = theta_key,
                                         matching_radius_arcsec = matching_radius_arcsec)
        
        print('Input images and catalogs saved to: ', self.input_images_path, self.input_catalogs_path)
        
    def register_reference_catalog(self, 
                                   catalog_path: Union[Path, str] = None,
                                   objname: str = None,
                                   list_ra: List[float] = None,
                                   list_dec: List[float] = None,
                                   gaiaxp_catalog_dir: Union[Path, str] = '/lyman/data1/factory/ref_cat'):
        reference_catalog_path = None
        if catalog_path is not None:
            if Path(catalog_path).exists():
                reference_catalog_path = str(catalog_path)
        if objname is not None:
            catalog_path = glob.glob(os.path.join(gaiaxp_catalog_dir, f'*{objname}*'))[0]
            if Path(catalog_path).exists():
                reference_catalog_path = str(catalog_path)
        if list_ra is not None and list_dec is not None:
            catalog_path = self.find_gaiaxp_catalog(list_ra = list_ra, list_dec = list_dec, gaiaxp_catalog_dir = gaiaxp_catalog_dir)
            if Path(catalog_path).exists():
                reference_catalog_path = str(catalog_path)
        
        if reference_catalog_path is None:
            raise RuntimeError('No reference catalog provided. Please provide a catalog path or objname or list_ra and list_dec.')
        self.configuration.inputs.gaiaxp_synphot_csv = reference_catalog_path                                   
        
    def find_gaiaxp_catalog(self, 
                            list_ra: List[float],
                            list_dec: List[float],
                            tileinfo_path: Union[Path, str] = '../tileinfo/7-DT/final_tiles.txt',
                            gaiaxp_catalog_dir: Union[Path, str] = '/lyman/data1/factory/ref_cat'):
        tiles = Tiles(tileinfo_path)
        matched_tile_tbl, matched_coords_dict, _ = tiles.find_overlapping_tiles(list_ra = list_ra, list_dec = list_dec, visualize = False)

        len_per_tile = dict()
        for tile_id, matched_coords in matched_coords_dict.items():
            len_per_tile[tile_id] = len(matched_coords)
        most_probable_tile = max(len_per_tile, key=len_per_tile.get)
        
        gaiaxp_catalog_path = glob.glob(os.path.join(gaiaxp_catalog_dir, f'*{most_probable_tile}*'))[0]
        return gaiaxp_catalog_path
    
    def run(self):
        original_dir = os.getcwd()
        os.chdir(self.workdir)
        self.configuration.to_yaml(path = self.configuration_path)
        
        # Check whether input_images and input_catalogs are saved
        if not self.input_images_saved or not self.input_catalogs_saved:
            raise ValueError('Input images and catalogs are not saved. Please register targets first. Run self.register_target() first.')
        if not self.reference_catalog_saved:
            raise ValueError('Reference catalog is not saved. Please register reference catalog first. Run self.register_reference_catalog() first.')
        
        # Use subprocess
        command = f"tract7dt run --config {self.configuration_path}"
        print('Running tract7dt...')
        try:
            subprocess.run(command, shell=True, check=True)
            
        except subprocess.CalledProcessError as e:
            print(f"Error during tract7dt execution: {e.stderr.decode()}")
            raise e
        finally:
            os.chdir(original_dir)
        
        
#%%

# %%
if __name__ == '__main__':
    from bridge.utils import HostGalaxyCatalog
    hg = HostGalaxyCatalog()
#%%
if __name__ == '__main__':
    from ezphot.utils import DataBrowser
    dbrowser = DataBrowser('scidata')
    dbrowser.objname = 'UDS'
    target_imgset = dbrowser.search(pattern = '*.fits', return_type = 'science')
    target_imgset.select_images(filter = ['m400','m425','m450','m475','m500','m525','m550','m575','m600','m625','m650','m675','m700','m725','m750','m775','m800','m825','m850','m875'])
    target_imgset.target_images.sort(key = lambda x: x.filter)
    image_paths = [img.path for img in target_imgset.target_images]
    catalog_paths = [img.catalog.path for img in target_imgset.target_images]
    filter_list = [img.filter for img in target_imgset.target_images]
#%%
# ================================================
# ZP CALIBRATION
# ================================================
#%% FIND REFERENCE CATALOG
if __name__ == '__main__':
    target_refcatset = dbrowser.search(pattern = '*.fits.refcat', return_type = 'catalog')
    for refcat in target_refcatset.catalogs:
        refcat.data
        print(refcat.nsources)
    merged_tbl, merged_metadata = target_refcatset.merge_catalogs(join_type = 'outer')       
# %% RUN TRACT7DT
if __name__ == '__main__':
    target_ra = []
    target_dec = []
    target_type = []
    # target_ra = [30.7663010688] # SN2025afih
    # target_dec = [4.24367725701]
    # target_ra = [101.796646]
    # target_dec = [-10.420567]
    # target_type = ['STAR']
    # host_tbl = hg.search_catalog(ra = target_ra[0], dec = target_dec[0], radius_deg = 10/3600)
    # if len(host_tbl) > 0:
    #     host_ra = host_tbl['RA'][0]
    #     host_dec = host_tbl['Dec'][0]
    #     host_type = 'EXP'
    #     target_ra.append(host_ra)
    #     target_dec.append(host_dec)
    #     target_type.append(host_type)        
#%%
if __name__ == '__main__':
    import numpy as np
    self = Tract7DTRunner(image_paths = image_paths, filter_list = filter_list, catalog_paths = catalog_paths)
    list_ra = list(merged_tbl['ra_basis']) 
    list_dec = list(merged_tbl['dec_basis']) 
    list_type = ['STAR'] * len(list_ra)
    list_ra.extend(target_ra)
    list_dec.extend(target_dec)
    list_type.extend(target_type)
    self.configuration.epsf.psf_model = 'gaussian'
    self.register_target(list_ra = list_ra, list_dec = list_dec, list_type = list_type, update_type_from_catalog = False, update_flux_from_catalog = False)
    self.register_reference_catalog(objname = dbrowser.objname)
    self.run()
# %% READ TRACT7DT RESULT
if __name__ == '__main__':
    from astropy.io import ascii
    result = ascii.read(self.workdir / 'final_catalog_with_fit.csv')
    flux_keys = [key for key in result.colnames if key.startswith('FLUX_') and not key.endswith('_fit')]
    filter_keys = []
    for flux_key in flux_keys:
        mag_key = flux_key.replace('FLUX_', 'MAG_')
        flux_key_fit = flux_key + '_fit'
        fluxerr_key_fit = flux_key.replace('FLUX_', 'FLUXERR_') + '_fit'
        mag_key_fit = mag_key + '_fit'
        magerr_key_fit = mag_key_fit.replace('MAG_', 'MAGERR_')
        filter = flux_key.replace('FLUX_', '')
        filter_keys.append(filter)
        flux_input = result[flux_key]
        flux_fit = result[flux_key_fit]    
        fluxerr_fit = result[fluxerr_key_fit]
        mag_fit = -2.5*np.log10(flux_fit)
        magerr_fit = 2.5/np.log(10) * fluxerr_fit / flux_fit
        result[mag_key_fit] = mag_fit
        result[magerr_key_fit] = magerr_fit
#%% MATCH PSF AND REF CATALOGS
if __name__ == '__main__':
    from ezphot.helper import Helper
    from astropy.coordinates import SkyCoord
    from ezphot.skycatalog import SkyCatalog
    refcatalog = SkyCatalog(objname = dbrowser.objname)
    ref_tbl = refcatalog.data
    result.sort('MAGERR_m400_fit')

    max_distance_arcsec = 10
    helper = Helper()
    ra_psf = result['RA']
    dec_psf = result['DEC']
    skycoord_psf = SkyCoord(ra_psf, dec_psf, unit = 'deg')
    ra_ref = ref_tbl['ra']
    dec_ref = ref_tbl['dec']
    skycoord_ref = SkyCoord(ra_ref, dec_ref, unit = 'deg')
    idx_psf, idx_ref, _ = helper.cross_match(skycoord_psf, skycoord_ref, max_distance_second = max_distance_arcsec)
    result_matched = result[idx_psf]#[:30]
    ref_tbl_matched = ref_tbl[idx_ref]#[:30]
#%% CALCULATE ZP FOR EACH FILTER
if __name__ == '__main__':
    helper = Helper()
    zp_all = dict()
    mag_psf_all = dict()
    magerr_psf_all = dict()
    mag_ref_all = dict()
    for flux_key in flux_keys:
        filter = flux_key.replace('FLUX_', '')
        
        fluxerr_key_fit = flux_key.replace('FLUX_', 'FLUXERR_') + '_fit'
        mag_key_fit = flux_key.replace('FLUX_', 'MAG_') + '_fit'
        magerr_key_fit = mag_key_fit.replace('MAG_', 'MAGERR_')
        mag_key_ref = f'{filter}_mag'
        
        mag_from_psf = result_matched[mag_key_fit]
        magerr_from_psf = result_matched[magerr_key_fit]
        mag_from_ref = ref_tbl_matched[mag_key_ref]
        zp_all[filter] = mag_from_ref - mag_from_psf
        mag_psf_all[filter] = mag_from_psf
        magerr_psf_all[filter] = magerr_from_psf
        mag_ref_all[filter] = mag_from_ref
#%% PLOT ZP FOR EACH FILTER
from astropy.stats import sigma_clipped_stats
if __name__ == '__main__':
    import matplotlib.pyplot as plt
    from ezphot.dataobjects import LightCurve
    zp_filter = dict()
    zp_filter_std = dict()
    color_map = LightCurve.FILTER_COLOR
    plt.figure(figsize = (10, 6))
    for filter_ in filter_list:
        mag_psf = mag_psf_all[filter_]
        mag_ref = mag_ref_all[filter_]
        zp = zp_all[filter_]
        
        zp_mean, zp_median, zp_std = sigma_clipped_stats(zp, sigma = 3, maxiters = 5)
        if filter_ == 'm550':
            plt.scatter(mag_psf + zp_median, zp, edgecolor = color_map[filter_], facecolor = 'none', marker = 'D', label = f'{filter_}[ZP = {zp_median:.2f} ± {zp_std:.2f}]')
            plt.axhline(zp_median, color = color_map[filter_], ls = '--', lw = 1)
        # plt.axhline(zp_median + zp_std, color = 'k', ls = '--', lw = 1)
        # plt.axhline(zp_median - zp_std, color = 'k', ls = '--', lw = 1)
        # plt.text(0.05, 0.90, f'ZP = {zp_median:.3f} ± {zp_std:.3f}', transform = plt.gca().transAxes, fontsize = 14)
        # plt.ylim(zp_median - 0.5, zp_median + 0.5)
        # plt.show()    
        zp_filter[filter_] = zp_median
        zp_filter_std[filter_] = zp_std
    plt.legend(ncols = 3)
    plt.ylabel(r'ZP [M$_{PSF}$ - M$_{GAIAXP}$]', fontsize = 14)
    plt.xlabel('Magnitude [mag]', fontsize = 14)
    #plt.ylim(24, 26)
    plt.show()
    
#%%

#%% COMPARE WITH EZPHOT RESULT
if __name__ == '__main__':
    catalogset = dbrowser.search(pattern = 'coadd**.fits.cat', return_type = 'catalog')
    catalogset.select_sources(target_ra[0], target_dec[0])
    from ezphot.dataobjects import PhotometricSpectrum
    photspec = PhotometricSpectrum(catalogset)
    flux_key_ezphot = 'MAGSKY_APER_2'
    fluxerr_key_ezphot = 'MAGERR_APER_2'
    zperr_key_ezphot = 'ZPERR_APER_2'
    photspec_result = photspec.plot(ra = target_ra[0], dec = target_dec[0], flux_key = flux_key_ezphot, fluxerr_key = fluxerr_key_ezphot)
    photspec_fig = list(photspec_result[0].values())[0]
    photspec_tbl = photspec_result[3]
    photspec_ax = list(photspec_result[2].values())[0]
# %%
if __name__ == '__main__':
    target_coord = SkyCoord(target_ra[0], target_dec[0], unit = 'deg')
    result_target = result[SkyCoord(result['RA'], result['DEC'], unit = 'deg').separation(target_coord).value < 10/3600]
    magdiff_all = []
    for flux_key in flux_keys:
        filter_key = flux_key.replace('FLUX_', '')
        fluxerr_key = flux_key.replace('FLUX_', 'FLUXERR_')
        flux_key_fit = flux_key + '_fit'
        fluxerr_key_fit = fluxerr_key + '_fit'
        filter = flux_key.replace('FLUX_', '')
        
        flux_fit = result_target[flux_key_fit]
        fluxerr_fit = result_target[fluxerr_key_fit]
        abmag_fit = -2.5*np.log10(flux_fit) + zp_filter[filter]
        abmagerr_fit = 2.5/np.log(10) * fluxerr_fit / flux_fit
        zperr_fit = zp_filter_std[filter]
        magerr_fit = np.sqrt(abmagerr_fit**2 + zperr_fit**2)
        
        row_photspec = photspec_tbl[photspec_tbl['filter'] == filter_key]
        abmag_photspec = row_photspec[flux_key_ezphot]
        magerr_photspec = np.sqrt(row_photspec[fluxerr_key_ezphot]**2 + row_photspec[zperr_key_ezphot]**2)
        
        wl = photspec.FILTER_PIVOT_WAVELENGTH_NM[filter_key]
        mag_diff = float(abmag_fit - abmag_photspec[0])
        photspec_ax.scatter(wl, abmag_fit, edgecolor = 'red', facecolor = 'none', marker = 'D')
        photspec_ax.errorbar(wl, abmag_fit, yerr = magerr_fit, color = 'r', fmt = 'none')
        photspec_ax.text(wl, 15.0, f'{mag_diff:.3f}', color = 'r', rotation = 90)
        magdiff_all.append(mag_diff)

#%%
if __name__ == '__main__':
    print(np.mean(magdiff_all), np.std(magdiff_all), np.median(magdiff_all))
#%%
photspec_fig
#%%
# # ================================================
# # Photometry
# # ================================================
# # RUN TRACT7DT (TARGET)
# if __name__ == '__main__':
#     import numpy as np
#     self = Tract7DTRunner(image_paths = image_paths, filter_list = filter_list, catalog_paths = catalog_paths)
#     list_ra = [233.857430764]
#     list_dec = [12.0577222937]
#     list_type = ['STAR'] * len(list_ra)
#     self.register_target(list_ra = list_ra, list_dec = list_dec, list_type = list_type, update_type_from_catalog = False)
#     self.register_reference_catalog(objname = 'T22956')
#     self.run()

# #%% READ TRACT7DT RESULT
# if __name__ == '__main__':
#     from astropy.io import ascii
#     result = ascii.read(self.workdir / 'final_catalog_with_fit.csv')    
# #%% COMPARE WITH EZPHOT RESULT
# if __name__ == '__main__':
#     catalogset = dbrowser.search(pattern = 'coadd*20250520*.fits.cat', return_type = 'catalog')
#     catalogset.select_sources(233.857430764, 12.0577222937)
#     from ezphot.dataobjects import PhotometricSpectrum
#     photspec = PhotometricSpectrum(catalogset)
#     flux_key_ezphot = 'MAGSKY_APER_2'
#     fluxerr_key_ezphot = 'MAGERR_APER_2'
#     zperr_key_ezphot = 'ZPERR_APER_2'
#     photspec_result = photspec.plot(ra = 233.857430764, dec = 12.0577222937, flux_key = flux_key_ezphot, fluxerr_key = fluxerr_key_ezphot)
#     photspec_fig = photspec_result[0]['2025-05-20 02:05']
#     photspec_tbl = photspec_result[3]
#     photspec_ax = photspec_result[2]['2025-05-20 02:05']
# # %%
# if __name__ == '__main__':
#     magdiff_all = []
#     chisq_all = []
#     for flux_key in flux_keys:
#         filter_key = flux_key.replace('FLUX_', '')
#         flux_key_fit = flux_key + '_fit'
#         filter = flux_key.replace('FLUX_', '')
#         flux_fit = result[flux_key_fit]
#         abmag_fit = -2.5*np.log10(flux_fit) + zp_filter[filter]
#         row_photspec = photspec_tbl[photspec_tbl['filter'] == filter_key]
#         abmag_photspec = row_photspec[flux_key_ezphot]
#         magerr_photspec = np.sqrt(row_photspec[fluxerr_key_ezphot]**2 + row_photspec[zperr_key_ezphot]**2)
#         chisq = (abmag_fit[0] - abmag_photspec[0])**2 / magerr_photspec[0]**2
#         wl = photspec.FILTER_PIVOT_WAVELENGTH_NM[filter_key]
#         mag_diff = float(abmag_fit[0] - abmag_photspec[0])
#         photspec_ax.scatter(wl, abmag_fit[0], edgecolor = 'blue', facecolor = 'none', marker = 'D')
#         photspec_ax.text(wl, 14.3, f'{mag_diff:.3f}', color = 'blue', rotation = 90)
#         magdiff_all.append(mag_diff)
#         chisq_all.append(chisq)