#!/usr/bin/env python
from dataclasses import dataclass, field
from typing import List, Tuple
from netCDF4 import Dataset, num2date, chartostring
from datetime import datetime, timedelta

from functools import wraps

from mapies import MAPIES
from util.func_tools import timeit, get_file_list, time_converter, inflate_array
import time
import numpy as np
import pandas as pd
import xarray as xr


class TROPOMI(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.lat=lat
        self.lon=lon
        #self.quality_flag_limit

        # Add quality value filter number


        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
        """
        module_dir = os.path.dirname(__file__)
        if config_file:
            try:
                config = yaml.safe_load(open(config_file))
            except FileNotFoundError:
                logging.error("This is not a config file")
        else:
            config = yaml.safe_load(open(os.path.join(module_dir, "config/satellite_config.yaml")))

        variable_dict = config[self.datatype]["variables"]
        self.time_var = variable_dict["time_variable"]
        self.lon_var = variable_dict["lon_variable"]
        self.lat_var = variable_dict["lat_variable"]

        # If obs_var passed then run data analysis on that variable if not pull default
        obs_var = kwargs.get("obs_var")
        if obs_var:
            self.obs_var = obs_var
        else:
            self.obs_var = variable_dict["obs_variable"]


    @timeit
    def preprocess_vars(self):
        """
        Preprocessing of the dataset
        """
        super().preprocess_vars()


        # no2 column values
        obs_dims = self.ds[self.obs_var].dims
        obs_shape = self.ds[self.obs_var].shape
        obs_attrs = self.ds[self.obs_var].attrs

        # Duplicate values in the time array to get the same shape as the flattened variable array
        self.time_values = inflate_array(self.time_values, self.time_shape, no2_shape)
        self.time_values_index = inflate_array(self.time_values_index, self.time_shape, obs_shape)
        # Lon and lat values
        lon_dims = self.ds[self.lon_var].dims
        lon_shape = self.ds[self.lon_var].shape
        lon_attrs = self.ds[self.lon_var].attrs
        self.lon_values = self.ds[self.lon_var].values.flatten()

        # Reindex lon column values
        self.lon_values = self.reindex(
            self.time_values_index, 
            self.lon_values,
            )


        lat_dims = self.ds[self.lat_var].dims
        lat_shape = self.ds[self.lat_var].shape
        lat_attrs = self.ds[self.lat_var].attrs
        self.lat_values = self.ds[self.lat_var].values.flatten()

        # Reindex aod column values
        self.lat_values = self.reindex(
            self.time_values_index, 
            self.lat_values,
            )

        # TODO: Check if this works well, if the inflated values are along the right dimensions
        self.obs = self.ds[self.obs_var].values.flatten()

        # Reindex no2 column values
        self.obs = self.reindex(
            self.time_values_index,
            self.obs, 
            )


    @timeit
    def read_nc(self):
        """
        Read netcdf files with xarray
        """

        files = get_file_list("/esarchive/obs/sentinel/tropomi/original_files/tropomi_OFFL_NO2/S5P_OFFL_L2__NO2____20230101T220249_20230101*.nc")
        print(files)


        # Open dataset with xarray and dask
        # Tropomi requires you choose the group
        # If we want to read anything other than a single tropomi dataset we will need to do it in parallel
        self.ds = xr.open_mfdataset(files, group="PRODUCT")


    @timeit
    def to_plumes(self):
        """
        Restructing the data to be passed to calculate the plumes 
        """

        coords = coords = dict(time=time_values,)

        ds = self.to_xarray(coords=coords, data_vars=data_vars)



if __name__ == "__main__":
    start_date = "202301011430"
    end_date = "202401011531"
    c = TROPOMI(start_date, end_date)
    c.read_nc()
    c.preprocess_vars()