"""
Module spatial referencing for flopy model objects

"""
import numpy as np


class StructuredSpatialReference:
    """
    a simple class to locate the model grid in x-y space

    Parameters
    ----------

    delr : numpy ndarray
        the model discretization delr vector

    delc : numpy ndarray
        the model discretization delc vector

    lenuni : int
        the length units flag from the discretization package

    xul : float
        the x coordinate of the upper left corner of the grid

    yul : float
        the y coordinate of the upper left corner of the grid

    rotation : float
        the counter-clockwise rotation (in degrees) of the grid

    proj4_str: str
        a PROJ4 string that identifies the grid in space. warning: case
        sensitive!

    Attributes
    ----------
    xedge : ndarray
        array of column edges

    yedge : ndarray
        array of row edges

    xgrid : ndarray
        numpy meshgrid of xedges

    ygrid : ndarray
        numpy meshgrid of yedges

    xcenter : ndarray
        array of column centers

    ycenter : ndarray
        array of row centers

    xcentergrid : ndarray
        numpy meshgrid of column centers

    ycentergrid : ndarray
        numpy meshgrid of row centers

    Notes
    -----

    xul and yul can be explicitly (re)set after SpatialReference
    instantiation, but only before any of the other attributes and methods are
    accessed

    """

    def __init__(
        self,
        delr=1.0,
        delc=1.0,
        lenuni=1,
        nlay=1,
        xul=None,
        yul=None,
        rotation=0.0,
        proj4_str=None,
        **kwargs,
    ):
        self.delc = np.atleast_1d(np.array(delc))
        self.delr = np.atleast_1d(np.array(delr))
        self.nlay = nlay
        self.lenuni = lenuni
        self.proj4_str = proj4_str
        self._reset()
        self.set_spatialreference(xul, yul, rotation)

    @classmethod
    def from_namfile_header(cls, namefile):
        # check for reference info in the nam file header
        header = []
        with open(namefile, "r") as f:
            for line in f:
                if not line.startswith("#"):
                    break
                header.extend(line.strip().replace("#", "").split(","))

        xul, yul = None, None
        rotation = 0.0
        proj4_str = None
        start_datetime = "1/1/1970"

        for item in header:
            if "xul" in item.lower():
                try:
                    xul = float(item.split(":")[1])
                except:
                    pass
            elif "yul" in item.lower():
                try:
                    yul = float(item.split(":")[1])
                except:
                    pass
            elif "rotation" in item.lower():
                try:
                    rotation = float(item.split(":")[1])
                except:
                    pass
            elif "proj4_str" in item.lower():
                try:
                    proj4_str = ":".join(item.split(":")[1:]).strip()
                except:
                    pass
            elif "start" in item.lower():
                try:
                    start_datetime = item.split(":")[1].strip()
                except:
                    pass

        return (
            cls(xul=xul, yul=yul, rotation=rotation, proj4_str=proj4_str),
            start_datetime,
        )

    def __setattr__(self, key, value):
        reset = True
        if key == "delr":
            super().__setattr__("delr", np.atleast_1d(np.array(value)))
        elif key == "delc":
            super().__setattr__("delc", np.atleast_1d(np.array(value)))
        elif key == "xul":
            super().__setattr__("xul", float(value))
        elif key == "yul":
            super().__setattr__("yul", float(value))
        elif key == "rotation":
            super().__setattr__("rotation", float(value))
        elif key == "lenuni":
            super().__setattr__("lenuni", int(value))
        elif key == "nlay":
            super().__setattr__("nlay", int(value))
        else:
            super().__setattr__(key, value)
            reset = False
        if reset:
            self._reset()

    def reset(self, **kwargs):
        for key, value in kwargs.items():
            setattr(self, key, value)

    def _reset(self):
        self._xgrid = None
        self._ygrid = None
        self._ycentergrid = None
        self._xcentergrid = None

    @property
    def nrow(self):
        return self.delc.shape[0]

    @property
    def ncol(self):
        return self.delr.shape[0]

    def __eq__(self, other):
        if not isinstance(other, StructuredSpatialReference):
            return False
        if other.xul != self.xul:
            return False
        if other.yul != self.yul:
            return False
        if other.rotation != self.rotation:
            return False
        if other.proj4_str != self.proj4_str:
            return False
        return True

    @classmethod
    def from_gridspec(cls, gridspec_file, lenuni=0):
        f = open(gridspec_file, "r")
        raw = f.readline().strip().split()
        nrow = int(raw[0])
        ncol = int(raw[1])
        raw = f.readline().strip().split()
        xul, yul, rot = float(raw[0]), float(raw[1]), float(raw[2])
        delr = []
        j = 0
        while j < ncol:
            raw = f.readline().strip().split()
            for r in raw:
                if "*" in r:
                    rraw = r.split("*")
                    for n in range(int(rraw[0])):
                        delr.append(float(rraw[1]))
                        j += 1
                else:
                    delr.append(float(r))
                    j += 1
        delc = []
        i = 0
        while i < nrow:
            raw = f.readline().strip().split()
            for r in raw:
                if "*" in r:
                    rraw = r.split("*")
                    for n in range(int(rraw[0])):
                        delc.append(float(rraw[1]))
                        i += 1
                else:
                    delc.append(float(r))
                    i += 1
        f.close()
        return cls(
            np.array(delr),
            np.array(delc),
            lenuni,
            xul=xul,
            yul=yul,
            rotation=rot,
        )

    @property
    def attribute_dict(self):
        return {
            "xul": self.xul,
            "yul": self.yul,
            "rotation": self.rotation,
            "proj4_str": self.proj4_str,
        }

    def set_spatialreference(self, xul=None, yul=None, rotation=0.0):
        """
        set spatial reference - can be called from model instance
        """

        # Set origin and rotation
        if xul is None:
            self.xul = 0.0
        else:
            self.xul = xul
        if yul is None:
            self.yul = np.add.reduce(self.delc)
        else:
            self.yul = yul
        self.rotation = rotation
        self._reset()

    def __repr__(self):
        s = f"xul:{self.xul:<G}, yul:{self.yul:<G}, rotation:{self.rotation:<G}, "
        s += f"proj4_str:{self.proj4_str}"
        return s

    @property
    def xedge(self):
        return self.get_xedge_array()

    @property
    def yedge(self):
        return self.get_yedge_array()

    @property
    def xgrid(self):
        if self._xgrid is None:
            self._set_xygrid()
        return self._xgrid

    @property
    def ygrid(self):
        if self._ygrid is None:
            self._set_xygrid()
        return self._ygrid

    @property
    def xcenter(self):
        return self.get_xcenter_array()

    @property
    def ycenter(self):
        return self.get_ycenter_array()

    @property
    def ycentergrid(self):
        if self._ycentergrid is None:
            self._set_xycentergrid()
        return self._ycentergrid

    @property
    def xcentergrid(self):
        if self._xcentergrid is None:
            self._set_xycentergrid()
        return self._xcentergrid

    def _set_xycentergrid(self):
        self._xcentergrid, self._ycentergrid = np.meshgrid(
            self.xcenter, self.ycenter
        )
        self._xcentergrid, self._ycentergrid = self.rotate(
            self._xcentergrid,
            self._ycentergrid,
            self.rotation,
            0,
            self.yedge[0],
        )
        self._xcentergrid += self.xul
        self._ycentergrid += self.yul - self.yedge[0]

    def _set_xygrid(self):
        self._xgrid, self._ygrid = np.meshgrid(self.xedge, self.yedge)
        self._xgrid, self._ygrid = self.rotate(
            self._xgrid, self._ygrid, self.rotation, 0, self.yedge[0]
        )
        self._xgrid += self.xul
        self._ygrid += self.yul - self.yedge[0]

    @staticmethod
    def rotate(x, y, theta, xorigin=0.0, yorigin=0.0):
        """
        Given x and y array-like values calculate the rotation about an
        arbitrary origin and then return the rotated coordinates.  theta is in
        degrees.

        """
        theta = -theta * np.pi / 180.0
        xrot = (
            xorigin
            + np.cos(theta) * (x - xorigin)
            - np.sin(theta) * (y - yorigin)
        )
        yrot = (
            yorigin
            + np.sin(theta) * (x - xorigin)
            + np.cos(theta) * (y - yorigin)
        )
        return xrot, yrot

    def get_extent(self):
        """
        Get the extent of the rotated and offset grid

        Return (xmin, xmax, ymin, ymax)

        """
        x0 = self.xedge[0]
        x1 = self.xedge[-1]
        y0 = self.yedge[0]
        y1 = self.yedge[-1]

        # upper left point
        x0r, y0r = self.rotate(x0, y0, self.rotation, 0, self.yedge[0])
        x0r += self.xul
        y0r += self.yul - self.yedge[0]

        # upper right point
        x1r, y1r = self.rotate(x1, y0, self.rotation, 0, self.yedge[0])
        x1r += self.xul
        y1r += self.yul - self.yedge[0]

        # lower right point
        x2r, y2r = self.rotate(x1, y1, self.rotation, 0, self.yedge[0])
        x2r += self.xul
        y2r += self.yul - self.yedge[0]

        # lower left point
        x3r, y3r = self.rotate(x0, y1, self.rotation, 0, self.yedge[0])
        x3r += self.xul
        y3r += self.yul - self.yedge[0]

        xmin = min(x0r, x1r, x2r, x3r)
        xmax = max(x0r, x1r, x2r, x3r)
        ymin = min(y0r, y1r, y2r, y3r)
        ymax = max(y0r, y1r, y2r, y3r)

        return xmin, xmax, ymin, ymax

    def get_grid_lines(self):
        """
        get the grid lines as a list
        """
        xmin = self.xedge[0]
        xmax = self.xedge[-1]
        ymin = self.yedge[-1]
        ymax = self.yedge[0]
        lines = []
        # Vertical lines
        for j in range(self.ncol + 1):
            x0 = self.xedge[j]
            x1 = x0
            y0 = ymin
            y1 = ymax
            x0r, y0r = self.rotate(x0, y0, self.rotation, 0, self.yedge[0])
            x0r += self.xul
            y0r += self.yul - self.yedge[0]
            x1r, y1r = self.rotate(x1, y1, self.rotation, 0, self.yedge[0])
            x1r += self.xul
            y1r += self.yul - self.yedge[0]
            lines.append([(x0r, y0r), (x1r, y1r)])

        # horizontal lines
        for i in range(self.nrow + 1):
            x0 = xmin
            x1 = xmax
            y0 = self.yedge[i]
            y1 = y0
            x0r, y0r = self.rotate(x0, y0, self.rotation, 0, self.yedge[0])
            x0r += self.xul
            y0r += self.yul - self.yedge[0]
            x1r, y1r = self.rotate(x1, y1, self.rotation, 0, self.yedge[0])
            x1r += self.xul
            y1r += self.yul - self.yedge[0]
            lines.append([(x0r, y0r), (x1r, y1r)])
        return lines

    def get_xcenter_array(self):
        """
        Return a numpy one-dimensional float array that has the cell center x
        coordinate for every column in the grid in model space - not offset
        or rotated.

        """
        x = np.add.accumulate(self.delr) - 0.5 * self.delr
        return x

    def get_ycenter_array(self):
        """
        Return a numpy one-dimensional float array that has the cell center x
        coordinate for every row in the grid in model space - not offset
        of rotated.

        """
        Ly = np.add.reduce(self.delc)
        y = Ly - (np.add.accumulate(self.delc) - 0.5 * self.delc)

        return y

    def get_xedge_array(self):
        """
        Return a numpy one-dimensional float array that has the cell edge x
        coordinates for every column in the grid in model space - not offset
        or rotated.  Array is of size (ncol + 1)

        """
        xedge = np.concatenate(([0.0], np.add.accumulate(self.delr)))
        return xedge

    def get_yedge_array(self):
        """
        Return a numpy one-dimensional float array that has the cell edge y
        coordinates for every row in the grid in model space - not offset or
        rotated. Array is of size (nrow + 1)

        """
        length_y = np.add.reduce(self.delc)
        yedge = np.concatenate(
            ([length_y], length_y - np.add.accumulate(self.delc))
        )
        return yedge

    def write_gridSpec(self, filename):
        """write a PEST-style grid specification file"""
        f = open(filename, "w")
        f.write(f"{self.delc.shape[0]:10d} {self.delr.shape[0]:10d}\n")
        f.write(f"{self.xul:15.6E} {self.yul:15.6E} {self.rotation:15.6E}\n")
        for c in self.delc:
            f.write(f"{c:15.6E} ")
        f.write("\n")
        for r in self.delr:
            f.write(f"{r:15.6E} ")
        f.write("\n")
        return

    def get_vertices(self, i, j):
        pts = []
        xgrid, ygrid = self.xgrid, self.ygrid
        pts.append([xgrid[i, j], ygrid[i, j]])
        pts.append([xgrid[i + 1, j], ygrid[i + 1, j]])
        pts.append([xgrid[i + 1, j + 1], ygrid[i + 1, j + 1]])
        pts.append([xgrid[i, j + 1], ygrid[i, j + 1]])
        pts.append([xgrid[i, j], ygrid[i, j]])
        return pts

    def interpolate(self, a, xi, method="nearest"):
        """
        Use the griddata method to interpolate values from an array onto the
        points defined in xi.  For any values outside of the grid, use
        'nearest' to find a value for them.

        Parameters
        ----------
        a : numpy.ndarray
            array to interpolate from.  It must be of size nrow, ncol
        xi : numpy.ndarray
            array containing x and y point coordinates of size (npts, 2). xi
            also works with broadcasting so that if a is a 2d array, then
            xi can be passed in as (xgrid, ygrid).
        method : {'linear', 'nearest', 'cubic'}
            method to use for interpolation (default is 'nearest')

        Returns
        -------
        b : numpy.ndarray
            array of size (npts)

        """
        try:
            from scipy.interpolate import griddata
        except:
            print("scipy not installed\ntry pip install scipy")
            return None

        # Create a 2d array of points for the grid centers
        points = np.empty((self.ncol * self.nrow, 2))
        points[:, 0] = self.xcentergrid.flatten()
        points[:, 1] = self.ycentergrid.flatten()

        # Use the griddata function to interpolate to the xi points
        b = griddata(points, a.flatten(), xi, method=method, fill_value=np.nan)

        # if method is linear or cubic, then replace nan's with a value
        # interpolated using nearest
        if method != "nearest":
            bn = griddata(points, a.flatten(), xi, method="nearest")
            idx = np.isnan(b)
            b[idx] = bn[idx]

        return b


class VertexSpatialReference:
    """
    a simple class to locate the model grid in x-y space

    Parameters
    ----------
    xvdict: dictionary
        dictionary of x-vertices {1: (0,1,1,0)}
    yvdict: dictionary
        dictionary of y-vertices {1: (1,0,1,0)}
    lenuni : int
        the length units flag from the discretization package
    xadj : float
        the x coordinate of the upper left corner of the grid
    yadj : float
        the y coordinate of the upper left corner of the grid
    rotation : float
        the counter-clockwise rotation (in degrees) of the grid
    proj4_str: str
        a PROJ4 string that identifies the grid in space. warning:
        case sensitive!
    Attributes
    ----------
    xedge : ndarray
        array of column edges
    yedge : ndarray
        array of row edges
    xgrid : ndarray
        numpy meshgrid of xedges
    ygrid : ndarray
        numpy meshgrid of yedges
    xcenter : ndarray
        array of column centers
    ycenter : ndarray
        array of row centers
    xcentergrid : ndarray
        numpy meshgrid of column centers
    ycentergrid : ndarray
        numpy meshgrid of row centers
    Notes
    -----

    xadj and yuadj can be explicitly (re)set after SpatialReference
    instantiation, but only before any of the other attributes and methods are
    accessed

    """

    def __init__(
        self,
        xvdict=None,
        yvdict=None,
        nlay=1,
        xadj=0,
        yadj=0,
        rotation=0.0,
        lenuni=1.0,
        proj4_str=None,
        **kwargs,
    ):

        assert len(xvdict) == len(
            yvdict
        ), f"len(xvdict): {len(xvdict)} != len(yvdict): {len(yvdict)}"

        self._xv = np.array([xvdict[idx] for idx, key in enumerate(xvdict)])
        self._yv = np.array([yvdict[idx] for idx, key in enumerate(yvdict)])
        self.nlay = nlay
        self.lenuni = lenuni
        self.proj4_str = proj4_str
        self._reset()
        self.set_spatialreference(xadj, yadj, rotation)

    def _reset(self):
        self._xvarr = None
        self._yvarr = None
        self._xvdict = None
        self._yvdict = None
        self._xyvdict = None
        self._xcenter_array = None
        self._ycenter_array = None
        self._ncpl = None

    @classmethod
    def from_namfile_header(cls, namefile):
        # check for reference info in the nam file header
        header = []
        with open(namefile, "r") as f:
            for line in f:
                if not line.startswith("#"):
                    break
                header.extend(line.strip().replace("#", "").split(","))

        xadj, yadj = None, None
        proj4_str = None
        start_datetime = "1/1/1970"

        for item in header:
            if "xadj" in item.lower():
                try:
                    xadj = float(item.split(":")[1])
                except:
                    pass
            elif "yadj" in item.lower():
                try:
                    yadj = float(item.split(":")[1])
                except:
                    pass
            elif "proj4_str" in item.lower():
                try:
                    proj4_str = ":".join(item.split(":")[1:]).strip()
                except:
                    pass
            elif "start" in item.lower():
                try:
                    start_datetime = item.split(":")[1].strip()
                except:
                    pass

        return cls(xdaj=xadj, yadj=yadj, proj4_str=proj4_str), start_datetime

    def __setattr__(self, key, value):
        reset = True
        if key == "xvdict":
            super().__setattr__("xvdict", dict(value))
        elif key == "yvdict":
            super().__setattr__("yvdict", dict(value))
        elif key == "xyvdict":
            super().__setattr__("xyvdict", value)
        elif key == "xadj":
            super().__setattr__("xadj", float(value))
        elif key == "yadj":
            super().__setattr__("yadj", float(value))
        elif key == "rotation":
            super().__setattr__("rotation", float(value))
        elif key == "lenuni":
            super().__setattr__("lenuni", int(value))
        else:
            super().__setattr__(key, value)
            reset = False
        if reset:
            self._reset()

    def set_spatialreference(self, xadj=0.0, yadj=0.0, rotation=0.0):
        """
        set spatial reference - can be called from model instance
        xadj, yadj should be named xadj, yadj since they represent an
        adjustment factor
        """

        # Set origin and rotation
        self.xadj = xadj
        self.yadj = yadj
        self.rotation = rotation
        self._reset()

    @staticmethod
    def rotate(x, y, theta, xorigin=0.0, yorigin=0.0):
        """
        Given x and y array-like values calculate the rotation about an
        arbitrary origin and then return the rotated coordinates.  theta is in
        degrees.

        """
        theta = -theta * np.pi / 180.0
        xrot = (
            xorigin
            + np.cos(theta) * (x - xorigin)
            - np.sin(theta) * (y - yorigin)
        )
        yrot = (
            yorigin
            + np.sin(theta) * (x - xorigin)
            + np.cos(theta) * (y - yorigin)
        )
        return xrot, yrot

    def get_extent(self):
        """
        Get the extent of the rotated and offset grid

        Return (xmin, xmax, ymin, ymax)

        """
        xvarr, yvarr = self._get_rotated_vertices()

        xmin = np.min(xvarr)
        xmax = np.max(xvarr)
        ymin = np.min(yvarr)
        ymax = np.max(yvarr)
        return xmin, xmax, ymin, ymax

    @property
    def ncpl(self):
        if self._ncpl is None:
            self._ncpl = self.xarr.size / self.nlay
        return self._ncpl

    @property
    def xdict(self):
        if self._xvdict is None:
            self._xvdict = self._set_vertices()[2]
        return self._xvdict

    @property
    def ydict(self):
        if self._yvdict is None:
            self._yvdict = self._set_vertices()[3]
        return self._yvdict

    @property
    def xydict(self):
        if self._xyvdict is None:
            self._xyvdict = self._set_vertices()[4]
        return self._xyvdict

    @property
    def xarr(self):
        if self._xvarr is None:
            self._xvarr = self._set_vertices()[0]
        return self._xvarr

    @property
    def yarr(self):
        if self._yvarr is None:
            self._yvarr = self._set_vertices()[1]
        return self._yvarr

    @property
    def xcenter_array(self):
        if self._xcenter_array is None:
            self._set_xcenter_array()
        return self._xcenter_array

    @property
    def ycenter_array(self):
        if self._ycenter_array is None:
            self._set_ycenter_array()
        return self._ycenter_array

    def _get_rotated_vertices(self):
        """
        Adjusts position and rotates vertices if applicable

        Returns
        -------
        xvarr, yvarr:
            rotated and adjusted np.arrays of vertices
        """
        xvadj = self._xv + self.xadj
        yvadj = self._yv + self.yadj
        xvarr, yvarr = self.rotate(xvadj, yvadj, self.rotation)

        return xvarr, yvarr

    def _set_vertices(self):
        """
        Sets variables _xvarr, _yvarr, _xvdict and _yvdict to be accessed by
        property instances.

        Returns
        -------

        """
        self._xvarr, self._yvarr = self._get_rotated_vertices()
        self._xvdict = {idx: verts for idx, verts in enumerate(self._xvarr)}
        self._yvdict = {idx: verts for idx, verts in enumerate(self._yvarr)}
        self._xyvdict = {
            idx: np.array(list(zip(self._xvdict[idx], self._yvdict[idx])))
            for idx in self._xvdict
        }
        return (
            self._xvarr,
            self._yvarr,
            self._xvdict,
            self._yvdict,
            self._xyvdict,
        )

    def _set_xcenter_array(self):
        """
        Gets the x vertex center location of all cells sets to a 1d array for
        further interpolation. Useful when using Scipy.griddata to contour data

        """
        if self._xvarr is None:
            self._set_vertices()

        self._xcenter_array = np.array([])
        for cell in self.xarr:
            self._xcenter_array = np.append(self._xcenter_array, np.mean(cell))

    def _set_ycenter_array(self):
        """
        Gets the x vertex center location of all cells sets to a 1d array for
        further interpolation. Useful when using Scipy.griddata to contour data

        Returns
        -------
        """

        if self._yvarr is None:
            self._set_vertices()

        self._ycenter_array = np.array([])
        for cell in self.yarr:
            self._ycenter_array = np.append(self._ycenter_array, np.mean(cell))


class SpatialReference:
    """
    A dynamic inheritance class that locates a gridded model in space

        Parameters
        ----------
        delr : numpy ndarray
            the model discretization delr vector
        delc : numpy ndarray
            the model discretization delc vector
        lenuni : int
            the length units flag from the discretization package
        xul : float
            the x coordinate of the upper left corner of the grid
        yul : float
           the y coordinate of the upper left corner of the grid
        rotation : float
            the counter-clockwise rotation (in degrees) of the grid
        proj4_str: str
            a PROJ4 string that identifies the grid in space. warning:
            case sensitive!
        xadj : float
            vertex grid: x vertex adjustment factor
        yadj : float
            vertex grid: y vertex adjustment factor
        xvdict: dict
            dictionary of x-vertices by cellnum ex. {0: (0,1,1,0)}
        yvdict: dict
            dictionary of y-vertices by cellnum ex. {0: (1,1,0,0)}
        distype: str
            model grid discretization type
    """

    def __new__(
        cls,
        delr=1.0,
        delc=1.0,
        xvdict=None,
        yvdict=None,
        lenuni=1,
        nlay=1,
        xul=None,
        yul=None,
        xadj=0.0,
        yadj=0.0,
        rotation=0.0,
        proj4_str=None,
        distype="structured",
    ):

        if distype == "structured":
            new = object.__new__(StructuredSpatialReference)
            new.__init__(
                delr=delr,
                delc=delc,
                lenuni=lenuni,
                nlay=nlay,
                xul=xul,
                yul=yul,
                rotation=rotation,
                proj4_str=proj4_str,
            )

        elif distype == "vertex":
            new = object.__new__(VertexSpatialReference)
            new.__init__(
                xvdict=xvdict,
                yvdict=yvdict,
                xadj=xadj,
                yadj=yadj,
                rotation=rotation,
                proj4_str=proj4_str,
                nlay=nlay,
            )

        elif distype == "unstructured":
            raise NotImplementedError(
                "Unstructured discretization not yet implemented"
            )

        else:
            raise TypeError(f"Discretization type {distype} not supported")

        return new
