# coding=utf-8
import netCDF4
import os

from earthdiagnostics.diagnostic import Diagnostic, DiagnosticBasinOption
from earthdiagnostics.utils import Utils, TempFile
# noinspection PyUnresolvedReferences
import earthdiagnostics.cdftoolspython as cdftoolspython
import numpy as np

from earthdiagnostics.modelingrealm import ModelingRealms
from earthdiagnostics.constants import Basins


class Siasiesiv(Diagnostic):
    """
    Compute the sea ice extent , area and volume  in both hemispheres or a specified region.


    :original author: Virginie Guemas <virginie.guemas@bsc.es>
    :contributor: Neven Fuckar <neven.fuckar@bsc.es>
    :contributor: Ruben Cruz <ruben.cruzgarcia@bsc.es>
    :contributor: Javier Vegas-Regidor <javier.vegas@bsc.es>

    :created: April 2012
    :last modified: June 2016

    """
    alias = 'siasiesiv'
    "Diagnostic alias for the configuration file"

    e1t = None
    e2t = None
    gphit = None

    def __init__(self, data_manager, startdate, member, chunk, basin, mask):
        """
        :param data_manager: data management object
        :type data_manager: DataManager
        :param startdate: startdate
        :type startdate: str
        :param member: member number
        :type member: int
        :param chunk: chunk's number
        :type chunk: int
        :param mask: mask to use
        :type mask: numpy.array
        """
        Diagnostic.__init__(self, data_manager)
        self.basin = basin
        self.startdate = startdate
        self.member = member
        self.chunk = chunk
        self.mask = mask
        self.required_vars = ['sit', 'sic']
        self.generated_vars = ['siextents', 'sivols', 'siareas', 'siextentn', 'sivoln', 'siarean']

    def __str__(self):
        return 'Siasiesiv Startdate: {0} Member: {1} Chunk: {2} Basin: {3}'.format(self.startdate, self.member,
                                                                                   self.chunk, self.basin.fullname)

    @classmethod
    def generate_jobs(cls, diags, options):
        """
        Creates a job for each chunk to compute the diagnostic

        :param diags: Diagnostics manager class
        :type diags: Diags
        :param options: basin
        :type options: list[str]
        :return:
        """
        options_available = (DiagnosticBasinOption('basin', Basins.Global), )
        options = cls.process_options(options, options_available)

        mask = Utils.get_mask(options['basin'])

        job_list = list()
        for startdate, member, chunk in diags.config.experiment.get_chunk_list():
            job_list.append(Siasiesiv(diags.data_manager, startdate, member, chunk, options['basin'], mask))
        mesh_handler = Utils.openCdf('mesh_hgr.nc')
        Siasiesiv.e1t = np.asfortranarray(mesh_handler.variables['e1t'][0, :])
        Siasiesiv.e2t = np.asfortranarray(mesh_handler.variables['e2t'][0, :])
        Siasiesiv.gphit = np.asfortranarray(mesh_handler.variables['gphit'][0, :])
        mesh_handler.close()

        return job_list

    def compute(self):
        """
        Runs the diagnostic
        """
        sit_file = self.data_manager.get_file(ModelingRealms.seaIce, 'sit', self.startdate, self.member, self.chunk)
        sit_handler = Utils.openCdf(sit_file)
        sit = np.asfortranarray(sit_handler.variables['sit'][:])
        timesteps = sit_handler.dimensions['time'].size
        sit_handler.close()

        sic_file = self.data_manager.get_file(ModelingRealms.seaIce, 'sic', self.startdate, self.member, self.chunk)
        sic_handler = Utils.openCdf(sic_file)
        Utils.convert_units(sic_handler.variables['sic'], '1.0')
        sic = np.asfortranarray(sic_handler.variables['sic'][:])
        sic_handler.close()

        result = np.empty((8, timesteps))
        for t in range(0, timesteps):
            try:

                result[:, t] = cdftoolspython.icediag.icediags(Siasiesiv.e1t, Siasiesiv.e2t, self.mask,
                                                               Siasiesiv.gphit, sit[t, :], sic[t, :])
            except Exception as ex:
                print ex

        self.send_file(self._extract_variable_and_rename(sit_file, result[4, :], 'sivols', '10^9 m3'),
                       ModelingRealms.seaIce, 'sivols', self.startdate, self.member, self.chunk,
                       region=self.basin.fullname)
        self.send_file(self._extract_variable_and_rename(sit_file, result[5, :], 'siareas', '10^9 m2'),
                       ModelingRealms.seaIce, 'siareas', self.startdate, self.member, self.chunk,
                       region=self.basin.fullname)
        self.send_file(self._extract_variable_and_rename(sit_file, result[7, :], 'siextents', '10^9 m2'),
                       ModelingRealms.seaIce, 'siextents', self.startdate, self.member, self.chunk,
                       region=self.basin.fullname)

        self.send_file(self._extract_variable_and_rename(sit_file, result[0, :], 'sivoln', '10^9 m3'),
                       ModelingRealms.seaIce, 'sivoln', self.startdate, self.member, self.chunk,
                       region=self.basin.fullname)
        self.send_file(self._extract_variable_and_rename(sit_file, result[1, :], 'siarean', '10^9 m2'),
                       ModelingRealms.seaIce, 'siarean', self.startdate, self.member, self.chunk,
                       region=self.basin.fullname)
        self.send_file(self._extract_variable_and_rename(sit_file, result[3, :], 'siextentn', '10^9 m2'),
                       ModelingRealms.seaIce, 'siextentn', self.startdate, self.member, self.chunk,
                       region=self.basin.fullname)

    @staticmethod
    def _extract_variable_and_rename(reference_file, values, cmor_name, units):
        temp = TempFile.get()
        reference_handler = Utils.openCdf(reference_file)
        os.remove(temp)
        handler = netCDF4.Dataset(temp, 'w')

        Utils.copy_variable(reference_handler, handler, 'time', add_dimensions=True)
        Utils.copy_variable(reference_handler, handler, 'time_bnds', add_dimensions=True, must_exist=False)
        Utils.copy_variable(reference_handler, handler, 'leadtime', add_dimensions=True, must_exist=False)
        reference_handler.close()

        new_var = handler.createVariable(cmor_name, float, 'time', fill_value=1.0e20)
        new_var.units = units
        new_var.short_name = cmor_name
        new_var.valid_min = 0.0
        new_var[:] = values
        new_var.valid_max = np.max(values)
        handler.close()
        return temp