#!/usr/bin/env python
from datetime import datetime, timedelta
from functools import wraps

import logging
import os
import time
#from tqdm import tqdm
import yaml
import multiprocessing
import gc
import dask.array as da
from glob import glob as glob_module
import sys

import numpy as np
import pandas as pd
import xarray as xr
import matplotlib.pyplot as plt
import cartopy
import cartopy.crs as ccrs
#from netCDF4 import Dataset

from .mapies import MAPIES
from mapies.util.func_tools import *
#timeit, get_file_list, time_converter, frequency_int, error_estimation, time_domain_selection, geo_to_rot
from mapies.grids.monarch import RotatedGrid, IrregularRotatedGrid, RegularGrid, regridding_function
from mapies.util.variables_names import VARS_ATTRIBUTES
from pathlib import Path
import json
from numpy.typing import DTypeLike, NDArray
from mapies.util.multiprocessing import VIIRSMultiprocessing

os.environ["OMP_NUM_THREADS"] = "8"

#logger = logging.getLogger(__name__)
#logging.basicConfig(filename='../logs/mapies.log', level=logging.INFO)


class VIIRS(MAPIES):
    """
    Class for VIIRS specific data
    """

    def __init__(self, start_date, end_date, **kwargs):
        """
        Inherited init class with new variables
        """
        super().__init__(start_date, end_date)
        
        self.time_orig = np.datetime64("1993-01-01T00:00:00")
        self.datatype="viirs"
        self.mp_type = VIIRSMultiprocessing
        self.dest = kwargs.get("dest")
        if not os.path.exists(self.dest):
            raise Exception("Output directory doesn't exist")
        self.indir = kwargs.get("indir")

        self.year = int(start_date[:4])
        self.start_day = int(start_date[6:8])  
        self.end_day = int(end_date[6:8]) 
        self.monthly_dict = {}
        self.daily_dict = {} # Dictionary to produce daily netcdf files 

        self.dates_slice = pd.date_range(
            self.start_date, 
            self.end_date, 
            ).strftime('%Y%m%d%H%M')

        if isinstance(self.start_date, str):
            self.start_date = time_converter(self.start_date)

        if isinstance(self.end_date, str):
            self.end_date = time_converter(self.end_date)

        self.read_config(**kwargs)


    @timeit
    def read_config(self, **kwargs):
        """
        Read yaml config file
        """
        super().read_config(**kwargs)
        #logger.info("Reading VIIRS specific config variables")

        da_dict = self.config[self.datatype]["da"]
        self.obsid_da = da_dict["obsid"] #int

        unc_const = kwargs.get("unc_const")
        if unc_const:
            self.unc_const = unc_const
        else:
            self.unc_const = da_dict["uncertainty_constants"] #float, int, list, dict
        


    @timeit
    def gather_nc_files(self):
        """
        Returns list of files needed to be processed by the application
        """
        file_patterns = []
        print(f' Date slice = {self.dates_slice}')
        if self.indir:
            for date in self.dates_slice:
                date = datetime.strptime(date, '%Y%m%d%H%M').strftime('%Y%j')

                filepaths = f'{self.indir}/**/AERDB_L2_VIIRS_NOAA20.A{date}*'
                file_patterns.append(filepaths) 

            files = sorted(get_file_list(file_patterns))
            print(f"Total number of files: {len(files)}")
        else:
            raise Exception("Please specify an input directory")
       
        files_filtered = filter_files(files, self.start_date, self.end_date)
        
        first_idx = max(0, files.index(files_filtered[0]) - 2)
        last_idx = min(len(files) - 1, files.index(files_filtered[-1]) + 2)
        self.files = files[first_idx : last_idx + 1]
        
        print(f"Total number of filtered files: {len(self.files)}")
        
        for file in self.files:
            parts = file.split('.')
            julian_day = parts[1][1:]
            date = datetime.strptime(julian_day, "%Y%j")
            if date.month not in self.monthly_dict:
                self.monthly_dict[date.month] = {
                "files": [],  # List to store file paths
                "lon": None,  # Longitudes array
                "lat": None,  # Latitudes array
                "obs": None,  # Observations array
                "count": None,  # Count of observations used to compute the average
                } 
            if date.day not in self.daily_dict:
                self.daily_dict[date.day] = {
                "files": [],  # List to store file paths
                "lon": None,  # Longitudes array
                "lat": None,  # Latitudes array
                "obs": None,  # Observations array
                "time": None, # Scan Start Time array
                }

            self.monthly_dict[date.month]["files"].append(file)
            self.daily_dict[date.day]["files"].append(file)
            
        print(f"Discovered files for months: {list(self.monthly_dict.keys())}")
        for month, data in self.monthly_dict.items():
            print(f"  Month {month:02d}: {len(data['files'])} files")
        
        print(f"Discovered files for days: {list(self.daily_dict.keys())}")
        
        for day, data in self.daily_dict.items():
            print(f"Day {day:02d}: {len(data['files'])} files")


# =============================================================================
# Supporting functions
# =============================================================================
    @timeit
    def to_da(self, frequency="D", save_nc= True, save_figs=False):
        """
        Function that returns all the needed variables for the DA
        """

        # Calculate error
        

        outfiles = []
        
        self.frequency = frequency_int(frequency)
        da_dates_slice = pd.date_range(
            self.start_date, 
            self.end_date,
            freq=frequency
            ).strftime('%Y%m%d%H%M')

        for date in da_dates_slice:
            l_border = (datetime.strptime(date, "%Y%m%d%H%M") - timedelta(hours=self.frequency/2)).strftime('%Y%m%d%H%M')
            r_border = (datetime.strptime(date, "%Y%m%d%H%M") + timedelta(hours=self.frequency/2)).strftime('%Y%m%d%H%M')

            filename = Path(self.dest).joinpath(f'obs{date}.nc')
            
            # concat all days that we need
            obs = np.array([])
            lon = np.array([])
            lat = np.array([])
            time_values = np.array([], dtype='datetime64')
            days = pd.date_range(l_border, r_border, freq="D")
            for pandas_day in days:
                day = pandas_day.day
                try:
                    day_obs = self.daily_dict[day]["obs"]
                    day_lon = self.daily_dict[day]["lon"]
                    day_lat = self.daily_dict[day]["lat"]
                    day_time = self.daily_dict[day]["time"]
                except KeyError:
                    print("No data for day selected")
                    continue
                obs = np.append(obs, day_obs)
                lon = np.append(lon, day_lon)
                lat = np.append(lat, day_lat)
                time_values = np.append(time_values, day_time)

            # Perform a further time cut
            try:
                obs, lon, lat, time_values = time_domain_selection(obs, lat, lon, time_values, l_border, r_border)
            except TypeError:
                print("No data, skipping")
                continue
            # Calculate obserr
            obserr = error_estimation(self.datatype, obs, self.unc_const)


            # Run aggregation with grid representation
            try:
                if self.grid_repr == "rotated":
                    lon, lat, rlon, rlat, obs, obserr, count_obs = regridding_function(self.grid, lon, lat, obs, obserr)
                elif self.grid_repr == "regular":
                    lon, lat, obs, obserr, count_obs = regridding_function(self.grid, lon, lat, obs, obserr)
            except ValueError:
                raise Exception("Have you specified a grid_repr")

            

            # Create new arrays with same length as obs
            obsid = np.full(shape=obs.shape, fill_value=self.obsid_da, dtype=self.int_type)
            obstype = np.full(shape=obs.shape, fill_value=30, dtype=self.int_type)
            obslev = np.full(shape=obs.shape, fill_value=-99999, dtype=self.int_type)
            obsn = np.full(shape=obs.shape, fill_value=count_obs, dtype=self.int_type)
            time_values = np.full(shape=obs.shape, fill_value=date, dtype=str)

            # Coords equals index

            assert obs.shape == obserr.shape
            assert lon.shape == time_values.shape
            coords = dict(index=("index", np.indices(obs.shape)[0, :]))
            data_vars = dict(
                time=(["index",], time_values),
                lon=(["index",], lon), 
                lat=(["index",], lat),
                obs=(["index",], obs),
                obserr=(["index", ], obserr),
                obsid=(["index", ], obsid),
                obstype=(["index", ], obstype),
                obslev=(["index", ], obslev),
                obsn=(["index", ], obsn),
            )
            if self.grid_repr == "rotated":
                rotated_dict = dict(
                    rlon=(["index",], rlon), 
                    rlat=(["index",], rlat), #Also return rlon/rlat values if available
                )
                data_vars.update(rotated_dict)

            
            #save the plots if specified for each time frequency
            if save_figs:
                self.plot_2D_obs_custom(lon, lat, obs, outdir=self.dest, filename=f"da_obs_{date}.png")

            ds = self.to_xarray(coords=coords, data_vars=data_vars)
            if save_nc:
                #logger.info(f"Outputting da data with {filename}")
                print(f"Outputting da data with {filename}")
                ds.to_netcdf(filename, encoding={})
            outfiles.append(filename)

        return outfiles