# -*- coding: utf-8 -*-
#************************************************************************
# --- Copyright (c) INRS 2016
# --- Institut National de la Recherche Scientifique (INRS)
# ---
# --- Distributed under the GNU Lesser General Public License, Version 3.0.
# --- See accompanying file LICENSE.txt.
#************************************************************************

from . import FEMesh
from . import DTData
try:
    from .DTReducOperation    import REDUCTION_OPERATION as Operation
except ImportError:
    from .DTReducOperation_pp import REDUCTION_OPERATION as Operation

import numpy as np

import logging
LOGGER = logging.getLogger("INRS.H2D2.Tools.Data.Field")

class DTField:
    def __init__(self):
        pass

    def getDataAtTime(self, reducOp = Operation.op_noop, tsteps = None, cols = []):
        raise NotImplementedError

    def getDataAtStep(self, reducOp = Operation.op_noop, tsteps = None, cols = []):
        raise NotImplementedError

    def getDataActual(self):
        raise NotImplementedError

    def getData(self):
        raise NotImplementedError

    def doInterpolate(self, X, Y):
        raise NotImplementedError

    def doProbe(self, X, Y):
        raise NotImplementedError

    def getGrid(self):
        raise NotImplementedError

    def getDataMin(self):
        raise NotImplementedError

    def getDataMax(self):
        raise NotImplementedError

class DT2DFiniteElementField(DTField):
    def __init__(self, mesh, data=None):
        super(DT2DFiniteElementField, self).__init__()
        self.mesh = mesh    # FEMesh or subgrid
        self.data = data    # DTData
        self.tstp = None    # (Operation, Time steps)
        self.cols = None    # Columns in data
        self.actu = None    # Activ data
        self.itrp = FEMesh.FEInterpolatedPoints(self.mesh) if self.mesh else None
        self.prbe = FEMesh.FEProbeOnPoints(self.mesh) if self.mesh else None
        self.ndgb = self.mesh.getNodeNumbersGlobal() if self.mesh and self.mesh.isSubMesh() else None

    def __copy__(self):
        """
        Shallow copy, while reseting the activ data
        """
        newone = type(self)(None) 
        newone.mesh = self.mesh
        newone.ndgb = self.ndgb
        newone.itrp = self.itrp
        newone.prbe = self.prbe
        newone.data = self.data
        newone.tstp = None
        newone.cols = None
        newone.actu = None
        return newone

    def __str__(self):
        if self.tstp:
            o, t = self.tstp
            if o == Operation.op_noop:
                s = ''
            else:
                s = '%s(' % o.name
            if len(t) == 1:
                s += '%f' % t[0]
            else:
                s += '[%f,%f]' % (t[0], t[-1])
            if o != Operation.op_noop:
                s += ')'
        else:
            s = ''
        return s

    def getDataAtTime(self, reducOp=Operation.op_noop, tsteps=None, cols=[]):
        LOGGER.debug('DT2DFiniteElementField: Get data with op %s at times (%s)', reducOp, tsteps)
        if (reducOp, tsteps) != self.tstp or cols != self.cols:
            op = DTData.DTDataReduc()
            op.compute(self.data, reducOp=reducOp, tsteps=tsteps, cols=cols)
            self.tstp = (reducOp, tsteps)
            self.cols = cols
            self.actu = op.getVal()[self.ndgb] if self.ndgb else op.getVal()
        return self.actu

    def getDataAtStep(self, reducOp=Operation.op_noop, tsteps=None, cols=[]):
        LOGGER.debug('DT2DFiniteElementField: Get data with op %s at steps (%s)', reducOp, tsteps)
        if (reducOp, tsteps) != self.tstp or cols != self.cols:
            op = DTData.DTDataReduc()
            op.compute(self.data, reducOp=reducOp, tsteps=tsteps, cols=cols)
            self.tstp = (reducOp, tsteps)
            self.cols = cols
            self.actu = op.getVal()[self.ndgb] if self.ndgb else op.getVal()
        return self.actu

    def getDataActual(self):
        return self.actu

    def getData(self):
        return self.data

    def doInterpolate(self, X, Y):
        self.itrp.setCoord(X, Y)
        i = self.itrp.interpolate(self.actu)
        return i

    def doProbe(self, X, Y):
        return self.prbe.interpolate(X, Y, self.actu)

    def getGrid(self):
        return self.mesh

    def getDataMin(self):
        if self.actu.ndim == 1:
            return np.min(self.actu)
        else:
            return np.min(np.linalg.norm(self.actu, axis=self.actu.ndim-1))

    def getDataMax(self):
        if self.actu.ndim == 1:
            return np.max(self.actu)
        else:
            return np.max(np.linalg.norm(self.actu, axis=self.actu.ndim-1))

class DT2DRegularGridField(DTField):
    def __init__(self, rgrid, source):
        super(DT2DRegularGridField, self).__init__()
        self.rgrid  = rgrid
        self.source = source    # DTField
        self.actu   = None      # Activ data

    def getDataAtTime(self, reducOp = Operation.op_noop, tsteps = None, cols = []):
        LOGGER.debug('DT2DRegularGridField: Get data with op %s at times (%s)', reducOp, tsteps)
        self.source.getDataAtTime(reducOp, tsteps, cols)
        self.actu = self.source.doInterpolate(*self.rgrid.getCoordinates())
        return self.actu

    def getDataAtStep(self, reducOp = Operation.op_noop, tsteps = None, cols = []):
        LOGGER.debug('DT2DRegularGridField: Get data with op %s at steps (%s)', reducOp, tsteps)
        self.source.getDataAtStep(reducOp, tsteps, cols)
        self.actu = self.source.doInterpolate(*self.rgrid.getCoordinates())
        return self.actu

    def getData(self):
        return self.source.getData()

    def getDataActual(self):
        #nx, ny = self.rgrid.getGridSize()
        #nd = np.size(self.actu[0])
        #shp = (nx, ny, nd)
        #if nd <= 1: shp = (nx, ny)
        #return np.array(self.actu).reshape(shp)
        return self.actu

    def doInterpolate(self, X, Y):
        raise NotImplementedError

    def doProbe(self, X, Y):
        raise NotImplementedError

    def getGrid(self):
        return self.rgrid

    def getDataMin(self):
        """
        return min ignoring nan
        """
        if self.actu.ndim == 1:
            return np.nanmin(self.actu)
        else:
            return np.nanmin(np.linalg.norm(self.actu, axis=self.actu.ndim-1))

    def getDataMax(self):
        """
        return max ignoring nan
        """
        if self.actu.ndim == 1:
            return np.nanmax(self.actu)
        else:
            return np.nanmax(np.linalg.norm(self.actu, axis=self.actu.ndim-1))

class DT1DRegularGridField(DTField):
    def __init__(self, rgrid, source):
        super(DT1DRegularGridField, self).__init__()
        self.rgrid  = rgrid
        self.source = source    # DTField
        self.actu   = None      # Activ data

    def getDataAtTime(self, reducOp = Operation.op_noop, tsteps = None, cols = []):
        LOGGER.debug('DT1DRegularGridField: Get data with op %s at times (%s)', reducOp, tsteps)
        self.source.getDataAtTime(reducOp, tsteps, cols)
        self.actu = self.source.doInterpolate(*self.rgrid.getCoordinates())
        return self.actu

    def getDataAtStep(self, reducOp = Operation.op_noop, tsteps = None, cols = []):
        LOGGER.debug('DT1DRegularGridField: Get data with op %s at steps (%s)', reducOp, tsteps)
        self.source.getDataAtStep(reducOp, tsteps, cols)
        self.actu = self.source.doInterpolate(*self.rgrid.getCoordinates())
        return self.actu

    def getData(self):
        return self.source.getData()

    def getDataActual(self):
        nx, ny = self.rgrid.getGridSize()
        nd = np.size(self.actu[0])
        shp = (nx, ny, nd)
        if nd <= 1: shp = (nx, ny)
        return np.array(self.actu).reshape(shp)

    def doInterpolate(self, X, Y):
        raise NotImplementedError

    def doProbe(self, X, Y):
        raise NotImplementedError

    def getGrid(self):
        return self.rgrid

    def getDataMin(self):
        """
        return min ignoring nan
        """
        return np.nanmin(self.actu)

    def getDataMax(self):
        """
        return max ignoring nan
        """
        return np.nanmax(self.actu)


if __name__ == "__main__":
    # pylint: disable=all
    def main():
        import os
        p = 'E:/Projets_simulation/EQ/Dry-Wet/Simulations/GLOBAL_01/Simulation/global01_0036'
        f = os.path.join(p, 'simul000.pst.sim')
        reader = DTData.constructReaderFromFile(f) #, cols = (3,))
        blocs = reader.readStructNoStats()

    streamHandler = logging.StreamHandler()
    LOGGER.addHandler(streamHandler)
    LOGGER.setLevel(logging.DEBUG)

    main()