# coding=utf-8
import os

# noinspection PyUnresolvedReferences
import earthdiagnostics.cdftoolspython as cdftoolspython
import netCDF4
import numpy as np
from bscearth.utils.log import Log

from earthdiagnostics.constants import Basins
from earthdiagnostics.diagnostic import Diagnostic, DiagnosticBasinOption
from earthdiagnostics.modelingrealm import ModelingRealms
from earthdiagnostics.utils import Utils, TempFile


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.generated = {}

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

    @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)

        if options['basin'] is None:
            Log.error('Basin not recognized')
            return ()

        mask = np.asfortranarray(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 request_data(self):
        self.sit = self.request_chunk(ModelingRealms.seaIce, 'sit', self.startdate, self.member, self.chunk)
        self.sic = self.request_chunk(ModelingRealms.seaIce, 'sic', self.startdate, self.member, self.chunk)

    def declare_data_generated(self):
        self._declare_var('sivols')
        self._declare_var('siareas')
        self._declare_var('siextents')

        self._declare_var('sivoln')
        self._declare_var('siarean')
        self._declare_var('siextentn')

    def _declare_var(self, var_name):
        self.generated[var_name] = self.declare_chunk(ModelingRealms.seaIce, var_name, self.startdate, self.member,
                                                      self.chunk, region=self.basin.name)

    def compute(self):
        """
        Runs the diagnostic
        """
        sit_handler = Utils.openCdf(self.sit.local_file)
        sit = np.asfortranarray(sit_handler.variables['sit'][:])
        timesteps = sit_handler.dimensions['time'].size
        sit_handler.close()

        sic_handler = Utils.openCdf(self.sic.local_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):
            result[:, t] = cdftoolspython.icediag.icediags(Siasiesiv.e1t, Siasiesiv.e2t, self.mask,
                                                           Siasiesiv.gphit, sit[t, :], sic[t, :])

        self._extract_variable_and_rename(result[4, :], 'sivols', '10^9 m3')
        self._extract_variable_and_rename(result[5, :], 'siareas', '10^9 m2')
        self._extract_variable_and_rename(result[7, :], 'siextents', '10^9 m2')

        self._extract_variable_and_rename(result[0, :], 'sivoln', '10^9 m3')
        self._extract_variable_and_rename(result[1, :], 'siarean', '10^9 m2')
        self._extract_variable_and_rename(result[3, :], 'siextentn', '10^9 m2')

    def _extract_variable_and_rename(self, values, cmor_name, units):
        temp = TempFile.get()
        reference_handler = Utils.openCdf(self.sit.local_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()
        self.generated[cmor_name].set_local_file(temp)