from parser import Parser
from autosubmit.config.log import Log
from earthdiagnostics.ocean import Salinity
from utils import Utils
from earthdiagnostics import cdftools
from autosubmit.date.chunk_date_lib import *
import shutil
import os


class Diags:
    def __init__(self, config_file):
        self._read_config(config_file)
        Utils.scratch_folder = os.path.join(self.scratch_dir)
        cdftools.path = self.cdftools_path

    def run(self):
        if not os.path.exists(self.scratch_dir):
            os.makedirs(self.scratch_dir)
        os.chdir(self.scratch_dir)

        # Copy mesh files
        Log.info('Copying mesh files')
        shutil.copy(os.path.join(self.con_files, 'mesh_mask_nemo.{0}.nc'.format(self.nemo_version)), 'mesh_hgr.nc')
        shutil.copy(os.path.join(self.con_files, 'mesh_mask_nemo.{0}.nc'.format(self.nemo_version)), 'mesh_zgr.nc')
        shutil.copy(os.path.join(self.con_files, 'mesh_mask_nemo.{0}.nc'.format(self.nemo_version)), 'mask.nc')
        shutil.copy(os.path.join(self.con_files, 'new_maskglo.{0}.nc'.format(self.nemo_version)), 'new_maskglo.nc')

        # dic_variables = dict()
        # dic_variables['x'] = 'i'
        # dic_variables['y'] = 'j'
        # dic_variables['z'] = 'lev'
        # dic_variables['nav_lon'] = 'lon'
        # dic_variables['nav_lat'] = 'lat'
        # dic_variables['nav_lev'] = 'lev'
        # dic_variables['time_counter'] = 'time'
        # dic_variables['t'] = 'time'
        #
        # Utils.rename_variables('mesh_hgr.nc', dic_variables, False, True)
        # Utils.rename_variables('mesh_zgr.nc', dic_variables, False, True)
        # Utils.rename_variables('mask.nc', dic_variables, False, True)

        if os.path.exists(os.path.join(self.con_files, 'mask.regions.{0}.nc'.format(self.nemo_version))):
            shutil.copy(os.path.join(self.con_files, 'mask.regions.{0}.nc'.format(self.nemo_version)),
                        'mask_regions.nc')

        # Check if cmorized and convert if not

        # Run diagnostics
        Log.info('Running diagnostics')
        for diag in self.diags.split():
            if diag == 'vert_mean_sal':
                Log.info("Running vert_mean_sal")
                min_depth = 0
                max_depth = 300
                for [input_file, output_file] in self._get_file_names('ocean', 'so',
                                                                      'vertmeansal'.format(min_depth, max_depth)):
                    Salinity.vertical_mean(str(input_file), str(output_file), min_depth, max_depth)
                Log.result('Finished vert_mean_sal')

        Utils.clean()

    def _read_config(self, config_file):
        self.parser = Parser()
        self.parser.optionxform = str
        self.parser.read(config_file)

        # Read diags config
        self.scratch_dir = self.parser.get_option('DIAGNOSTICS', 'SCRATCH_DIR')
        self.data_dir = self.parser.get_option('DIAGNOSTICS', 'DATA_DIR')
        self.con_files = self.parser.get_option('DIAGNOSTICS', 'CON_FILES')
        self.diags = self.parser.get_option('DIAGNOSTICS', 'DIAGS')
        self.frequency = self.parser.get_option('DIAGNOSTICS', 'FREQUENCY')
        self.cdftools_path = self.parser.get_option('DIAGNOSTICS', 'CDFTOOLS_PATH')

        # Read experiment config
        self.institute = self.parser.get_option('EXPERIMENT', 'INSTITUTE')
        self.expid = self.parser.get_option('EXPERIMENT', 'EXPID')
        self.members = self.parser.get_option('EXPERIMENT', 'MEMBERS')
        self.startdates = self.parser.get_option('EXPERIMENT', 'STARTDATES')
        self.chunk_size = self.parser.get_int_option('EXPERIMENT', 'CHUNK_SIZE')
        self.chunks = self.parser.get_int_option('EXPERIMENT', 'CHUNKS')
        self.model = self.parser.get_option('EXPERIMENT', 'MODEL')
        self.nemo_version = self.parser.get_option('EXPERIMENT', 'NEMO_VERSION')

        self.scratch_dir = os.path.join(self.scratch_dir, 'diags', self.expid)

    def _get_file_names(self, domain, input_var, output_var):
        file_names = list()

        for startdate in self.startdates.split():
            start = parse_date(startdate)
            for member in self.members:
                member_plus = str(int(member)+1)
                member_path = os.path.join(self.data_dir,self.expid, startdate, 'fc'+member, 'outputs', 'output',
                                           self.institute, self.model, self.model, 'S'+startdate, self.frequency,
                                           domain)
                for chunk in range(1, self.chunks + 1):
                    chunk_start = chunk_start_date(start, chunk, self.chunk_size, 'month', 'standard')
                    chunk_end = chunk_end_date(chunk_start, self.chunk_size, 'month', 'standard')
                    chunk_end = previous_day(chunk_end, 'standard')

                    input_file = os.path.join(member_path, input_var, 'r{0}i1p1'.format(member_plus),
                                              '{0}_{1[0]}{2}_{3}_output_r{4}i1p1_'
                                              '{5}-{6}.nc'.format(input_var, domain.upper(),
                                                                  self.frequency, self.model, member_plus,
                                                                  "{0:04}{1:02}".format(chunk_start.year,
                                                                                        chunk_start.month),
                                                                  "{0:04}{1:02}".format(chunk_end.year,
                                                                                        chunk_end.month)))
                    output_file = os.path.join(member_path, input_var, 'r{0}i1p1'.format(member_plus),
                                               '{0}_{1[0]}{2}_{3}_output_r{4}i1p1_'
                                               '{5}-{6}.nc'.format(output_var, domain.upper(),
                                                                   self.frequency, self.model, member_plus,
                                                                   "{0:04}{1:02}".format(chunk_start.year,
                                                                                         chunk_start.month),
                                                                   "{0:04}{1:02}".format(chunk_end.year,
                                                                                         chunk_end.month)))
                    file_names.append([input_file, output_file])

        return file_names


def main():
    Log.set_console_level(Log.DEBUG)
    diags = Diags('/home/Earth/jvegas/pyCharm/ocean_diagnostics/earthdiagnostics/diags.conf')
    diags.run()
    # Utils.clean()


if __name__ == "__main__":
    main()