import netCDF4
import os

from earthdiagnostics import Utils, cdftools, cdo, TempFile
from earthdiagnostics.basins import Basins
from autosubmit.config.log import Log
import shutil


class Ice(object):
    """
    Class containing all the diagnostics related to ice
    """

    @staticmethod
    def siasiesiv(thickness_file, fraction_file, basin, svolume_file, sarea_file, sextent_file,
                  nvolume_file, narea_file, nextent_file):
        """
        Compute the sea ice extent (1000km2), area (1000km2), volume (km3)
        and mean thickness (m) in both hemispheres or a specified region.

        Created in April 2012         Author : vguemas@ic3.cat
        Modified in June 2014         Author : neven.fuckar@ic3.cat

        Computation of the properties in various selected regions according to
        mask.regions.${NEMOVERSION}.nc (mask_regions.nc) is based on modification
        of mask.regions.ORCA1.noverticalinfo.Matt.nc from Matthieu Chevallier.
        :param svolume_file:
        :param sarea_file:
        :param sextent_file:
        :param nvolume_file:
        :param narea_file:
        :param nextent_file:
        :param fraction_file:
        :param thickness_file:
        :param output_file: output file name
        :param basin: region of interest
        :return:
        """
        nco = Utils.nco
        input_scratch = TempFile.get()
        shutil.copy(thickness_file, input_scratch)
        nco.ncks(input=fraction_file, output=input_scratch, options='-A -v sic')
        ntime = int(cdo.ntime(input=input_scratch)[0])
        files = list()

        basin = Basins.parse(basin)
        if basin != Basins.Global:
            shutil.move('mask.nc', 'original_mask.nc')
            shutil.move('mask_regions.nc', 'mask.nc')
            Utils.rename_variable('mask.nc', basin.fullname, 'tmask')
        error = None
        try:
            for time in range(ntime):
                Log.info('Running time {0}', time)
                temp = TempFile.get()
                nco.ncks(input=input_scratch, output=temp, options='-O -d time,{0}'.format(time))
                cdftools.run('cdficediags', input=temp)
                Utils.move_file('icediags.nc', temp)
                files.append(temp)
        except Exception as ex:
            error = ex.message
        finally:
            if basin != Basins.Global:
                Utils.rename_variable('mask.nc', 'tmask', basin.fullname)
                shutil.move('mask.nc', 'mask_regions.nc')
                shutil.move('original_mask.nc', 'mask.nc')
            if error:
                raise Exception(error)

        temp = TempFile.get()
        cdo.cat(input=' '.join(files), output=temp)
        nco.ncks(input=input_scratch, output=temp, options='-A -v time')

        Ice.extract_variable_and_rename(temp, 'SVolume', 'sivols', svolume_file, "10^3 km3")
        Ice.extract_variable_and_rename(temp, 'SArea', 'siareas', sarea_file, "10^6 km2")
        Ice.extract_variable_and_rename(temp, 'SExnsidc', 'siextents', sextent_file, "10^6 km2")

        Ice.extract_variable_and_rename(temp, 'NVolume', 'sivoln', nvolume_file, "10^3 km3")
        Ice.extract_variable_and_rename(temp, 'NArea', 'siarean', narea_file, "10^6 km2")
        Ice.extract_variable_and_rename(temp, 'NExnsidc', 'siextentn', nextent_file, "10^6 km2")

        os.remove(temp)

    @staticmethod
    def extract_variable_and_rename(input_file, variable, cmor_name, output_file, output_units):
        temp = TempFile.get()
        # Utils.nco.ncks(input=input_file, output=temp, options='-O -v {0},time,time_bnds'.format(variable))
        input_handler = Utils.cdo.openCdf(input_file)

        os.remove(temp)
        handler = netCDF4.Dataset(temp, 'w')

        # Create dimensions
        handler.createDimension('time', None)
        handler.createDimension('bnds', 2)

        # Copy time variable
        time_var = input_handler.variables['time']
        new_time = handler.createVariable('time', time_var.datatype, time_var.dimensions)
        new_time.setncatts({k: time_var.getncattr(k) for k in time_var.ncattrs()})
        new_time[:] = time_var[:]

        original_bnds = input_handler.variables['time_bnds']
        new_bnds = handler.createVariable('time_bnds', original_bnds.datatype, original_bnds.dimensions)
        new_bnds.setncatts({k: original_bnds.getncattr(k) for k in original_bnds.ncattrs()})
        new_bnds[:] = original_bnds[:]

        original_variable = input_handler.variables[variable]
        values = original_variable[:, 0, 0]

        new_var = handler.createVariable(cmor_name, original_variable.datatype, 'time')
        new_var.setncatts({k: original_variable.getncattr(k) for k in original_variable.ncattrs()})
        factor = Ice.get_conversion_factor(original_variable.units, output_units)
        values *= factor
        new_var[:] = values
        new_var.units = output_units
        new_var.short_name = cmor_name
        handler.close()

        Utils.move_file(temp, output_file)

    @staticmethod
    def get_conversion_factor(input_units, output_units):
        units = input_units.split()
        if len(units) == 1:
            scale_unit = 1
            unit = units[0]
        else:
            if '^' in units[0]:
                values = units[0].split('^')
                scale_unit = pow(int(values[0]), int(values[1]))
            else:
                scale_unit = float(units[0])
            unit = units[1]

        units = output_units.split()
        if len(units) == 1:
            scale_new_unit = 1
            new_unit = units[0]
        else:
            if '^' in units[0]:
                values = units[0].split('^')
                scale_new_unit = pow(int(values[0]), int(values[1]))
            else:
                scale_new_unit = float(units[0])
            new_unit = units[1]

        factor = Ice.get_factor(new_unit, unit)
        invert = False
        if factor is None:
            factor = Ice.get_factor(unit, new_unit)
            invert = True

        if factor is None:
            raise Exception("Conversion from {0} to {1} not supported".format(units, output_units))

        if invert:
            factor = scale_unit / float(scale_new_unit * factor)
        else:
            factor = (factor * scale_unit) / float(scale_new_unit)
        return factor

    @staticmethod
    def get_factor(new_unit, unit):
        # Add  only the conversions with a factor greater than 1
        if unit == new_unit:
            return 1
        if unit == 'km3':
            if new_unit == 'm3':
                return pow(1000, 3)
        elif unit == 'km2':
            if new_unit == 'm2':
                return pow(1000, 2)
        return None