import os
import glob
import xarray as xr
import numpy as np
import pandas as pd
from tqdm import tqdm
import multiprocessing as mp
import pathlib
from copy import deepcopy

from . import core, utils



class Timeseries:
    ''' Initialize a CESM Timeseries case generated by CESM Postprocessing

    Args:
        root_dir (str): the root directory of the CESM Timeseries output
        grid_dict (dict): the grid dictionary for different components
        timestep (int): the number of years stored in a single timeseries file
    '''
    def __init__(self, root_dir, grid_dict=None, casename=None):
        self.path_pattern='comp/proc/tseries/month_1/casename.mdl.h_str.vn.timespan.nc'
        self.root_dir = os.path.abspath(root_dir)
        self.casename = casename

        self.grid_dict = {'atm': 'ne30pg3', 'ocn': 'g16'}
        if grid_dict is not None:
            self.grid_dict.update(grid_dict)

        self.grid_dict['lnd'] = self.grid_dict['atm']
        self.grid_dict['rof'] = self.grid_dict['atm']
        self.grid_dict['ice'] = self.grid_dict['ocn']

        utils.p_header(f'>>> case.root_dir: {self.root_dir}')
        utils.p_header(f'>>> case.path_pattern: {self.path_pattern}')
        utils.p_header(f'>>> case.grid_dict: {self.grid_dict}')
        if self.casename is not None:
            utils.p_header(f'>>> case.casename: {self.casename}')

        self.paths = utils.find_paths(self.root_dir, self.path_pattern)

        self.ds = {}
        self.diags = {}
        self.vars_info = {}
        for path in self.paths:
            comp = path.split('/')[-5]
            mdl = path.split('.')[-5]
            h_str = path.split('.')[-4]
            vn = path.split('.')[-3]
            if (vn, comp) not in self.vars_info:
                self.vars_info[(vn, comp)] = (comp, mdl, h_str)

        utils.p_success(f'>>> case.vars_info created')

    def get_paths(self, vn, comp=None, timespan=None):
        if comp is None: comp = self.get_vn_comp(vn)
        comp, mdl, h_str = self.vars_info[(vn, comp)]
        paths = utils.find_paths(self.root_dir, self.path_pattern, vn=vn, comp=comp, mdl=mdl, h_str=h_str)
        if timespan is None:
            paths_sub = paths
        else:
            syr, eyr = timespan
            paths_sub = []
            for path in paths:
                syr_tmp = int(path.split('.')[-2].split('-')[0][:4])
                eyr_tmp = int(path.split('.')[-2].split('-')[1][:4])
                if (syr_tmp >= syr and syr_tmp <= eyr) or (eyr_tmp >= syr and eyr_tmp <= eyr):
                    paths_sub.append(path)

        return paths_sub

    def get_vn_comp(self, vn):
        comps = []
        for (v, comp) in self.vars_info:
            if v == vn:
                comps.append(comp)
        
        if len(comps) == 1:
            return comps[0]
        elif len(comps) == 0:
            if f'get_{vn}' in self.diags.DiagCalc.__dict__:
                utils.p_warning(f'>>> {vn} is a supported derived variable.')
            else:
                raise ValueError('The input variable name is unknown.')
        else:
            utils.p_warning(f'{vn} belongs to components: {comps}')
            raise ValueError('The input variable name belongs to multiple components. Please specify via the argument `comp`.')

    
    def load(self, vn, comp=None, timespan=None, load_idx=-1, adjust_month=True, verbose=True, **kws):
        if comp is None:
            comp = self.get_vn_comp(vn)

        if (vn, comp) in self.vars_info:
            if timespan is None:
                paths = self.get_paths(vn, comp=comp)[load_idx]
            else:
                paths = self.get_paths(vn, comp=comp, timespan=timespan)

            if vn in self.ds:
                if self.ds[vn].path != paths:
                    if verbose: utils.p_warning(f'>>> case.ds["{vn}"] will be reloaded due to different paths.')
                    self.clear_ds(vn)

            if vn not in self.ds:
                comp, mdl, h_str = self.vars_info[(vn, comp)]

                _kws = {
                   'vn': vn,
                   'adjust_month': adjust_month, 
                   'comp': comp,
                   'grid': self.grid_dict[comp],
                }
                _kws.update(kws)
                if not isinstance(paths, (list, tuple)):
                    ds =  core.open_dataset(paths, **_kws)
                else:
                    ds =  core.open_mfdataset(paths, **_kws)

                self.ds[vn] = ds
                self.ds[vn].attrs['vn'] = vn
                if verbose: utils.p_success(f'>>> case.ds["{vn}"] created')
            # else:
            #     if verbose: utils.p_warning(f'>>> case.ds["{vn}"] already loaded; to reload, run case.clear_ds("{vn}") before case.load("{vn}")')

        else:
            if verbose: utils.p_warning(f'>>> Variable {vn} not existing')


    def get_climo(self, vn, comp=None, timespan=None, adjust_month=True, slicing=False, regrid=False, dlat=1, dlon=1):
        ''' Generate the climatology file for the given variable

        Args:
            slicing (bool): could be problematic
        '''
        if comp is None: comp = self.get_vn_comp(vn)
        grid = self.grid_dict[comp]
        paths = self.get_paths(vn, comp=comp, timespan=timespan)
        ds = core.open_mfdataset(paths, adjust_month=adjust_month)

        if slicing: ds = ds.sel(time=slice(timespan[0], timespan[1]))
        ds_out = ds.x.climo
        ds_out.attrs['comp'] = comp
        ds_out.attrs['grid'] = grid
        if regrid: ds_out = ds_out.x.regrid(dlat=dlat, dlon=dlon)
        return ds_out

    def save_climo(self, output_dirpath, vn, comp=None, timespan=None, adjust_month=True,
                   slicing=False, regrid=False, dlat=1, dlon=1, overwrite=False):

        output_dirpath = pathlib.Path(output_dirpath)
        if not output_dirpath.exists():
            output_dirpath.mkdir(parents=True, exist_ok=True)
            utils.p_success(f'>>> output directory created at: {output_dirpath}')

        fname = f'{vn}_climo.nc' if self.casename is None else f'{self.casename}_{vn}_climo.nc'
        out_path = os.path.join(output_dirpath, fname)
        if overwrite or not os.path.exists(out_path):
            if os.path.exists(out_path): os.remove(out_path)
            if comp is None: comp = self.get_vn_comp(vn)

            climo = self.get_climo(
                vn, comp=comp, timespan=timespan, adjust_month=adjust_month,
                slicing=slicing, regrid=regrid, dlat=dlat, dlon=dlon,
            )
            climo.to_netcdf(out_path)
            climo.close()

    def gen_climo(self, output_dirpath, comp=None, timespan=None, vns=None, adjust_month=True,
                  nproc=1, slicing=False, regrid=False, dlat=1, dlon=1, overwrite=False):

        if comp is None:
            raise ValueError('Please specify component via the argument `comp`.')

        if vns is None:
            vns = [k[0] for k, v in self.vars_info.items() if v[0]==comp]

        utils.p_header(f'>>> Generating climo for {len(vns)} variables:')
        for i in range(len(vns)//10+1):
            print(vns[10*i:10*i+10])

        if nproc == 1:
            for v in tqdm(vns, total=len(vns), desc=f'Generating climo files'):
                self.save_climo(
                    output_dirpath, v, comp=comp, timespan=timespan,
                    adjust_month=adjust_month, slicing=slicing,
                    regrid=regrid, dlat=dlat, dlon=dlon,
                    overwrite=overwrite,
                )
        else:
            utils.p_hint(f'>>> nproc: {nproc}')
            with mp.Pool(processes=nproc) as p:
                arg_list = [(output_dirpath, v, comp, timespan, adjust_month, slicing, regrid, dlat, dlon, overwrite) for v in vns]
                p.starmap(self.save_climo, tqdm(arg_list, total=len(vns), desc=f'Generating climo files'))

        utils.p_success(f'>>> {len(vns)} climo files created in: {output_dirpath}')

    def get_ts(self, vn, comp, timespan=None, adjust_month=True, slicing=False, regrid=False, dlat=1, dlon=1):
        grid = self.grid_dict[comp]
        paths = self.get_paths(vn, comp=comp, timespan=timespan)
        ds = core.open_mfdataset(paths, adjust_month=adjust_month)

        if slicing: ds = ds.sel(time=slice(timespan[0], timespan[1]))

        ds_out = ds
        ds_out.attrs['comp'] = comp
        ds_out.attrs['grid'] = grid
        if regrid: ds_out = ds_out.x.regrid(dlat=dlat, dlon=dlon)
        return ds_out

    def save_means(self, vn, comp, output_dirpath, timespan, adjust_month=True, slicing=False, regrid=False, dlat=1, dlon=1, overwrite=False):
        output_dirpath = pathlib.Path(output_dirpath)
        if not output_dirpath.exists():
            output_dirpath.mkdir(parents=True, exist_ok=True)
            utils.p_success(f'>>> output directory created at: {output_dirpath}')

        ds = self.get_ts(vn, comp, timespan=timespan, adjust_month=adjust_month, slicing=slicing, regrid=False)

        sn_dict = {
            'ANN': list(range(1, 13)),
            'DJF': [12, 1, 2],
            'MAM': [1, 2, 3],
            'JJA': [6, 7, 8],
            'SON': [9, 10, 11],
        }

        for sn, months in sn_dict.items():
            output_subdirpath = pathlib.Path(os.path.join(output_dirpath, sn))
            if not output_subdirpath.exists():
                output_subdirpath.mkdir(parents=True, exist_ok=True)

            fname = f'{timespan[0]}_{timespan[1]}_{vn}_{sn}_means.nc'
            if self.casename is not None: fname = f'{self.casename}_{fname}'

            out_path = os.path.join(output_subdirpath, fname)
            if overwrite or not os.path.exists(out_path):
                ds_ann = ds.x.annualize(months=months)
                if regrid: ds_ann = ds_ann.x.regrid(dlat=dlat, dlon=dlon)
                ds_ann.to_netcdf(out_path)
                ds_ann.close()

    def gen_means(self, output_dirpath, comp=None, vns=None, timespan=None, adjust_month=True, slicing=False,
                  regrid=False, dlat=1, dlon=1, overwrite=False, nproc=1):

        if comp is None:
            raise ValueError('Please specify component via the argument `comp`.')

        if vns is None:
            vns = [k[0] for k, v in self.vars_info.items() if v[0]==comp]

        utils.p_header(f'>>> Generating seaonal means for {len(vns)} variables:')
        for i in range(len(vns)//10+1):
            print(vns[10*i:10*i+10])

        if nproc == 1:
            for vn in vns:
                self.save_means(
                    vn, comp, output_dirpath, timespan, adjust_month=adjust_month, slicing=slicing,
                    regrid=regrid, dlat=dlat, dlon=dlon, overwrite=overwrite, 
                )
        else:
            utils.p_hint(f'>>> nproc: {nproc}')
            with mp.Pool(processes=nproc) as p:
                arg_list = [(vn, comp, output_dirpath, timespan, adjust_month, slicing, regrid, dlat, dlon, overwrite) for vn in vns]
                p.starmap(self.save_means, tqdm(arg_list, total=len(vns), desc=f'Generating seasonal mean files'))


    def clear_ds(self, vn=None):
        ''' Clear the existing `.ds` property
        '''
        if vn is not None:
            try:
                self.ds.pop(vn)
            except:
                pass
        else:
            self.ds = {}
        
    def copy(self):
        return deepcopy(self)

class Climo:
    def __init__(self, root_dir, casename):
        self.root_dir = root_dir
        self.casename = casename
        utils.p_header(f'>>> case.root_dir: {self.root_dir}')
        utils.p_header(f'>>> case.casename: {self.casename}')

    def gen_MONS_climo(self, output_dirpath, climo_period=None):
        output_dirpath = pathlib.Path(output_dirpath)
        if not output_dirpath.exists():
            output_dirpath.mkdir(parents=True, exist_ok=True)
            utils.p_header(f'>>> output directory created at: {output_dirpath}')

        paths = sorted(glob.glob(os.path.join(self.root_dir, '*_climo.nc')))
        ds = core.open_mfdataset(paths)
        if climo_period is None:
            try:
                climo_period = ds.attrs['climo_period']
            except:
                pass

        if climo_period is not None:
            fname = f'{self.casename}_{climo_period[0]}_{climo_period[1]}_MONS_climo.nc'
        else:
            fname = f'{self.casename}_MONS_climo.nc'

        output_fpath = os.path.join(output_dirpath, fname)
        ds.to_netcdf(output_fpath)
        utils.p_header(f'>>> MONS_climo generated at: {output_fpath}')
        self.MONS_climo_path = output_fpath
        utils.p_success(f'>>> case.MONS_climo_path created')

    def gen_seasons_climo(self, MONS_climo_path=None):
        if MONS_climo_path is None:
            MONS_climo_path = self.MONS_climo_path

        ds = xr.open_dataset(MONS_climo_path)

        sn_dict = {
            'ANN': list(range(1, 13)),
            'DJF': [12, 1, 2],
            'MAM': [1, 2, 3],
            'JJA': [6, 7, 8],
            'SON': [9, 10, 11],
        }

        for sn, months in sn_dict.items():
            output_fpath = MONS_climo_path.replace('MONS', sn)
            if os.path.exists(output_fpath):
                os.remove(output_fpath)

            ds_mean = ds.sel(time=months).mean('time').expand_dims('time')
            ds_mean = ds_mean.assign_coords(time=[ds.coords['time'][0]])
            ds_mean.to_netcdf(output_fpath, unlimited_dims={'time':True})
            utils.p_header(f'>>> {sn}_climo generated at: {output_fpath}')

class Means:
    def __init__(self, root_dir):
        self.root_dir = root_dir
        utils.p_header(f'>>> case.root_dir: {self.root_dir}')

    def merge_means(self, sn, output_dirpath, overwrite=False, casetag=None):
        utils.p_header(f'>>> Processing season {sn}')
        paths = glob.glob(os.path.join(self.root_dir, sn, f'*_{sn}_means.nc'))
        print(paths)
        if casetag is None:
            fname = f'{sn}_means.nc'
        else:
            fname = f'{casetag}_{sn}_means.nc'
        out_path = os.path.join(output_dirpath, fname)
        if overwrite or not os.path.exists(out_path):
            ds_list = []
            for path in paths:
                ds_tmp = core.open_dataset(path)
                for k, v in ds_tmp.coords.items():
                    try:
                        if any(np.isnan(v.values)):
                            ds_tmp = ds_tmp.drop_vars(k)
                    except:
                        pass
                ds_list.append(ds_tmp)

            utils.p_header(f'>>> Merging files')
            ds = xr.merge(ds_list)
            ds.to_netcdf(out_path)
            utils.p_header(f'>>> Merged mean file saved at: {out_path}')
        else:
            utils.p_warning(f'>>> The result already exists. Skipping ...')

    def merge_means_nproc(self, output_dirpath, sns=['ANN', 'DJF', 'MAM', 'JJA', 'SON'], overwrite=False, casetag=None, nproc=1):
        if nproc == 1:
            for sn in sns:
                self.merge_means(sn, output_dirpath, overwrite=overwrite, casetag=casetag)
        else:
            utils.p_hint(f'>>> nproc: {nproc}')
            with mp.Pool(processes=nproc) as p:
                arg_list = [(sn, output_dirpath, overwrite, casetag) for sn in sns]
                p.starmap(self.merge_means, tqdm(arg_list, total=len(sns), desc=f'Merging mean files'))