# coding=utf-8
"""Extract levels from variable"""
from earthdiagnostics.box import Box
from earthdiagnostics.diagnostic import (
    Diagnostic,
    DiagnosticOption,
    DiagnosticDomainOption,
    DiagnosticVariableListOption,
    DiagnosticIntOption,
)
from earthdiagnostics.utils import Utils, TempFile


class SelectLevels(Diagnostic):
    """
    Extract levels from file

    Parameters
    ----------
    data_manager: DataManager
    startdate: str
    member: int
    chunk: init
    domain: ModellingRealm
    variable: str
    grid: str or None
    first_level: int
    last_level: int
    """

    alias = "selev"
    "Diagnostic alias for the configuration file"

    def __init__(
        self,
        data_manager,
        startdate,
        member,
        chunk,
        domain,
        variable,
        grid,
        first_level,
        last_level,
    ):
        Diagnostic.__init__(self, data_manager)
        self.startdate = startdate
        self.member = member
        self.chunk = chunk
        self.variable = variable
        self.domain = domain
        self.grid = grid
        self.box = Box(False)
        self.box.min_depth = first_level
        self.box.max_depth = last_level

        self.variable_file = None
        self.result = None

    def __str__(self):
        return (
            "Select levels Startdate: {0.startdate} Member: {0.member} "
            "Chunk: {0.chunk} Variable: {0.domain}:{0.variable} "
            "Levels: {0.box.min_depth}-{0.box.max_depth} "
            "Grid: {0.grid}".format(self)
        )

    def __eq__(self, other):
        if self._different_type(other):
            return False
        return (
            self.startdate == other.startdate
            and self.member == other.member
            and self.chunk == other.chunk
            and self.domain == other.domain
            and self.variable == other.variable
            and self.box == other.box
            and self.grid == self.grid
        )

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

        :param diags: Diagnostics manager class
        :type diags: Diags
        :param options: domain,variables,grid
        :type options: list[str]
        :return:
        """
        options_available = (
            DiagnosticDomainOption(),
            DiagnosticVariableListOption(
                diags.data_manager.config.var_manager, "variables"
            ),
            DiagnosticIntOption("first_level"),
            DiagnosticIntOption("last_level"),
            DiagnosticOption("grid", ""),
        )
        options = cls.process_options(options, options_available)
        job_list = list()
        variables = options["variables"]
        for var in variables:
            chunk_list = diags.config.experiment.get_chunk_list()
            for startdate, member, chunk in chunk_list:
                job_list.append(
                    SelectLevels(
                        diags.data_manager,
                        startdate,
                        member,
                        chunk,
                        options["domain"],
                        var,
                        options["grid"],
                        options["first_level"],
                        options["last_level"],
                    )
                )
        return job_list

    def request_data(self):
        """Request data required by the diagnostic"""
        self.variable_file = self.request_chunk(
            self.domain,
            self.variable,
            self.startdate,
            self.member,
            self.chunk,
            grid=self.grid,
            to_modify=True,
        )

    def declare_data_generated(self):
        """Request data required by the diagnostic"""
        self.result = self.declare_chunk(
            self.domain,
            self.variable,
            self.startdate,
            self.member,
            self.chunk,
            grid=self.grid,
        )

    def compute(self):
        """Run the diagnostic"""
        temp = TempFile.get()
        handler = Utils.open_cdf(self.variable_file.local_file)
        var_name = ""
        for var in ("lev", "plev"):
            if var in handler.variables:
                var_name = var
                continue
        handler.close()

        Utils.nco().ncks(
            input=self.variable_file.local_file,
            output=temp,
            options="-O -d {1},{0.min_depth},{0.max_depth}".format(
                self.box, var_name
            ),
        )
        self.result.set_local_file(temp)

    @staticmethod
    def _create_var(var_name, var_values, source, destiny):
        old_var = source.variables[var_name]
        new_var = destiny.createVariable(
            var_name, old_var.dtype, dimensions=(var_name,)
        )
        new_var[:] = var_values
        Utils.copy_attributes(new_var, old_var)

        vertices_name = "{0}_vertices".format(var_name)

        if vertices_name in source.variables:
            var_vertices = source.variables[vertices_name]
            if var_name == "lon":
                vertices_values = var_vertices[0:1, ...]
            else:
                vertices_values = var_vertices[:, 0:1, :]
            new_lat_vertices = destiny.createVariable(
                vertices_name,
                var_vertices.dtype,
                dimensions=(var_name, "vertices"),
            )
            new_lat_vertices[:] = vertices_values
            Utils.copy_attributes(new_lat_vertices, var_vertices)