import warnings

import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np
import properties
import scipy.sparse as sp
from discretize import TensorMesh
from discretize import utils as dis_utils
from matplotlib.colors import LogNorm
from scipy.spatial import cKDTree as kdtree
from simpeg import utils


def setup_mpl():
    # https://matplotlib.org/stable/tutorials/introductory/customizing.html#the-default-matplotlibrc-file
    mpl.rc("lines", linewidth=3)
    mpl.rc("axes", titlesize=12)
    mpl.rc("xtick", labelsize=8)
    mpl.rc("ytick", labelsize=8)


def set_mesh_1d(hz):
    return TensorMesh([hz], x0=[0])


class Stitched1DModel(properties.HasProperties):
    topography = properties.Array("topography (x, y, z)", dtype=float, shape=("*", "*"))

    physical_property = properties.Array("Physical property", dtype=float)

    # line = properties.Array("Line", dtype=float, default=None)

    time_stamp = properties.Array("Time stamp", dtype=float, default=None)

    hz = properties.Array("Vertical thickeness of 1D mesh", dtype=float)

    def __init__(self, **kwargs):
        super(Stitched1DModel, self).__init__(**kwargs)
        warnings.warn("code under construction - API might change in the future")

    @property
    def n_sounding(self):
        if getattr(self, "_n_sounding", None) is None:
            self._n_sounding = self.topography.shape[0]
        return self._n_sounding

    @property
    def unique_line(self):
        if getattr(self, "_unique_line", None) is None:
            if self.line is None:
                raise Exception("line information is required!")
            self._unique_line = np.unique(self.line)
        return self._unique_line

    @property
    def xyz(self):
        if getattr(self, "_xyz", None) is None:
            xyz = np.empty((self.hz.size, self.topography.shape[0], 3), order="F")
            for i_xy in range(self.topography.shape[0]):
                z = -self.mesh_1d.cell_centers_x + self.topography[i_xy, 2]
                x = np.ones_like(z) * self.topography[i_xy, 0]
                y = np.ones_like(z) * self.topography[i_xy, 1]
                xyz[:, i_xy, :] = np.c_[x, y, z]
            self._xyz = xyz
        return self._xyz

    @property
    def mesh_1d(self):
        if getattr(self, "_mesh_1d", None) is None:
            if self.hz is None:
                raise Exception("hz information is required!")
            self._mesh_1d = set_mesh_1d(np.r_[self.hz])
        return self._mesh_1d

    @property
    def mesh_3d(self):
        if getattr(self, "_mesh_3d", None) is None:
            print(">> Automatically generating a 3D mesh.")
            self._mesh_3d = self.get_3d_mesh()
        return self._mesh_3d

    @property
    def physical_property_matrix(self):
        if getattr(self, "_physical_property_matrix", None) is None:
            if self.physical_property is None:
                raise Exception("physical_property information is required!")
            self._physical_property_matrix = self.physical_property.reshape(
                (self.hz.size, self.n_sounding), order="F"
            )
        return self._physical_property_matrix

    @property
    def distance(self):
        if getattr(self, "_distance", None) is None:
            self._distance = np.zeros(self.n_sounding, dtype=float)
            for line_tmp in self.unique_line:
                ind_line = self.line == line_tmp
                xy_line = self.topography[ind_line, :2]
                distance_line = np.r_[
                    0, np.cumsum(np.sqrt((np.diff(xy_line, axis=0) ** 2).sum(axis=1)))
                ]
                self._distance[ind_line] = distance_line
        return self._distance

    def plot_section(
        self,
        i_layer=0,
        i_line=0,
        x_axis="x",
        show_layer=False,
        plot_type="contour",
        physical_property=None,
        clim=None,
        ax=None,
        cmap="viridis",
        ncontour=20,
        scale="log",
        show_colorbar=True,
        aspect=1,
        zlim=None,
        dx=20.0,
        invert_xaxis=False,
        alpha=0.7,
        pcolorOpts={},
    ):
        ind_line = self.line == self.unique_line[i_line]
        if physical_property is not None:
            physical_property_matrix = physical_property.reshape(
                (self.hz.size, self.n_sounding), order="F"
            )
        else:
            physical_property_matrix = self.physical_property_matrix

        if x_axis.lower() == "y":
            x_ind = 1
            xlabel = "Northing (m)"
        elif x_axis.lower() == "x":
            x_ind = 0
            xlabel = "Easting (m)"
        elif x_axis.lower() == "distance":
            xlabel = "Distance (m)"

        if ax is None:
            fig = plt.figure(figsize=(15, 10))
            ax = plt.subplot(111)

        if clim is None:
            vmin = np.percentile(physical_property_matrix, 5)
            vmax = np.percentile(physical_property_matrix, 95)
        else:
            vmin, vmax = clim

        if scale == "log":
            norm = LogNorm(vmin=vmin, vmax=vmax)
            vmin = None
            vmax = None
        else:
            norm = None

        ind_line = np.arange(ind_line.size)[ind_line]

        for i in ind_line:
            inds_temp = [i]
            if x_axis == "distance":
                x_tmp = self.distance[i]
            else:
                x_tmp = self.topography[i, x_ind]

            topo_temp = np.c_[x_tmp - dx, x_tmp + dx]
            out = ax.pcolormesh(
                topo_temp,
                -self.mesh_1d.nodes_x + self.topography[i, 2],
                physical_property_matrix[:, inds_temp],
                cmap=cmap,
                alpha=alpha,
                vmin=vmin,
                vmax=vmax,
                norm=norm,
                shading="auto",
                **pcolorOpts
            )

        if show_layer:
            ax.plot(
                x_tmp,
                self.topography[ind_line, 2] - self.mesh_1d.cell_centers_x[i_layer],
                "--",
                lw=1,
                color="grey",
            )

        if show_colorbar:
            from mpl_toolkits import axes_grid1

            cb = plt.colorbar(out, ax=ax, fraction=0.01)
            cb.set_label("Conductivity (S/m)")

        ax.set_aspect(aspect)
        ax.set_xlabel(xlabel)
        ax.set_ylabel("Elevation (m)")
        if zlim is not None:
            ax.set_ylim(zlim)

        if x_axis == "distance":
            xlim = (
                self.distance[ind_line].min() - dx,
                self.distance[ind_line].max() + dx,
            )
        else:
            xlim = (
                self.topography[ind_line, x_ind].min() - dx,
                self.topography[ind_line, x_ind].max() + dx,
            )
        if invert_xaxis:
            ax.set_xlim(xlim[1], xlim[0])
        else:
            ax.set_xlim(xlim)

        plt.tight_layout()

        if show_colorbar:
            return out, ax, cb
        else:
            return out, ax
        return ax

    def get_3d_mesh(
        self,
        dx=None,
        dy=None,
        dz=None,
        npad_x=0,
        npad_y=0,
        npad_z=0,
        core_z_length=None,
        nx=100,
        ny=100,
    ):
        xmin, xmax = self.topography[:, 0].min(), self.topography[:, 0].max()
        ymin, ymax = self.topography[:, 1].min(), self.topography[:, 1].max()
        zmin, zmax = self.topography[:, 2].min(), self.topography[:, 2].max()
        zmin -= self.mesh_1d.nodes_x.max()

        lx = xmax - xmin
        ly = ymax - ymin
        lz = zmax - zmin

        if dx is None:
            dx = lx / nx
            print((">> dx:%.1e") % (dx))
        if dy is None:
            dy = ly / ny
            print((">> dy:%.1e") % (dy))
        if dz is None:
            dz = np.median(self.mesh_1d.hx)

        nx = int(np.floor(lx / dx))
        ny = int(np.floor(ly / dy))
        nz = int(np.floor(lz / dz))

        if nx * ny * nz > 1e6:
            warnings.warn(
                ("Size of the mesh (%i) will greater than 1e6") % (nx * ny * nz)
            )
        hx = [(dx, npad_x, -1.2), (dx, nx), (dx, npad_x, -1.2)]
        hy = [(dy, npad_y, -1.2), (dy, ny), (dy, npad_y, -1.2)]
        hz = [(dz, npad_z, -1.2), (dz, nz)]

        zmin = self.topography[:, 2].max() - dis_utils.unpack_widths(hz).sum()
        self._mesh_3d = TensorMesh([hx, hy, hz], x0=[xmin, ymin, zmin])

        return self.mesh_3d

    @property
    def P(self):
        if getattr(self, "_P", None) is None:
            raise Exception("Run get_interpolation_matrix first!")
        return self._P

    def get_interpolation_matrix(self, npts=20, epsilon=None):
        tree_2d = kdtree(self.topography[:, :2])
        xy = utils.ndgrid(self.mesh_3d.cell_centers_x, self.mesh_3d.cell_centers_y)

        distance, inds = tree_2d.query(xy, k=npts)
        if epsilon is None:
            epsilon = np.min([self.mesh_3d.h[0].min(), self.mesh_3d.h[1].min()])

        w = 1.0 / (distance + epsilon) ** 2
        w = utils.sdiag(1.0 / np.sum(w, axis=1)) * (w)
        I = utils.mkvc(np.arange(inds.shape[0]).reshape([-1, 1]).repeat(npts, axis=1))
        J = utils.mkvc(inds)

        self._P = sp.coo_matrix(
            (utils.mkvc(w), (I, J)), shape=(inds.shape[0], self.topography.shape[0])
        )

        mesh_1d = TensorMesh([np.r_[self.hz[:-1], 1e20]])

        z = self.P * self.topography[:, 2]

        self._actinds = dis_utils.active_from_xyz(self.mesh_3d, np.c_[xy, z])

        Z = np.empty(self.mesh_3d.vnC, dtype=float, order="F")
        Z = self.mesh_3d.gridCC[:, 2].reshape(
            (self.mesh_3d.shape_cells[0] * self.mesh_3d.shape_cells[1], self.mesh_3d.shape_cells[2]), order="F"
        )
        ACTIND = self._actinds.reshape(
            (self.mesh_3d.shape_cells[0] * self.mesh_3d.shape_cells[1], self.mesh_3d.shape_cells[2]), order="F"
        )

        self._Pz = []

        # This part can be cythonized or parallelized
        for i_xy in range(self.mesh_3d.shape_cells[0] * self.mesh_3d.shape_cells[1]):
            actind_temp = ACTIND[i_xy, :]
            z_temp = -(Z[i_xy, :] - z[i_xy])
            self._Pz.append(mesh_1d.get_interpolation_matrix(z_temp[actind_temp]))

    def interpolate_from_1d_to_3d(self, physical_property_1d):
        physical_property_2d = self.P * (
            physical_property_1d.reshape((self.hz.size, self.n_sounding), order="F").T
        )
        physical_property_3d = (
            np.ones(
                (self.mesh_3d.shape_cells[0] * self.mesh_3d.shape_cells[1], self.mesh_3d.shape_cells[2]),
                order="C",
                dtype=float,
            )
            * np.nan
        )

        ACTIND = self._actinds.reshape(
            (self.mesh_3d.shape_cells[0] * self.mesh_3d.shape_cells[1], self.mesh_3d.shape_cells[2]), order="F"
        )

        for i_xy in range(self.mesh_3d.shape_cells[0] * self.mesh_3d.shape_cells[1]):
            actind_temp = ACTIND[i_xy, :]
            physical_property_3d[i_xy, actind_temp] = (
                self._Pz[i_xy] * physical_property_2d[i_xy, :]
            )

        return physical_property_3d


def generate_average_map(mesh, values, z_top, z_bottom, averaging_method="arithmetic"):
    """
    Calculate vertical average of values in the 3D grid.

    Parameters
    ----------

    mesh: object
        Discretize TensorMesh object
    values: (*,) array_like, float
        3D coarse fraction model
    z_top: float
        top of the integration interval
    z_bottom: float
        bottom of the integration interval
    averaging_method: string
        option for average method (arithematic, harmonic, logarithmic)
    Returns
    -------
    averaged_values : (*,2) ndarray, float
        averaged values.
    """
    nx, ny, nz = mesh.vnC
    # this assumes the uniform cell size
    dz = mesh.hz[0]
    z_values = mesh.cell_centers_z
    tmp_values = values.copy().reshape((nx * ny, nz), order="F")
    inds_tmp = np.logical_and(z_values > z_bottom, z_values < z_top)
    tmp_values[:, ~inds_tmp] = np.nan
    inds_nan = np.isnan(tmp_values)
    tmp_values[inds_nan] = 0.0
    thickness = (~inds_nan).astype(float).sum(axis=1) * dz
    if averaging_method == "arithmetic":
        averaged_values = (tmp_values * dz).sum(axis=1) / thickness
    elif averaging_method == "harmonic":
        averaged_values = 1.0 / ((1.0 / tmp_values * dz).sum(axis=1) / thickness)
    elif averaging_method == "logarithmic":
        averaged_values = np.exp((np.log(tmp_values) * dz).sum(axis=1) / thickness)
    return averaged_values


def plot_layer_model(
    sig, mesh, xscale="log", ax=None, showlayers=False, xlim=None, **kwargs
):
    """
    Plot Conductivity model for the layered earth model
    """
    with warnings.catch_warnings():
        warnings.filterwarnings("ignore", category=DeprecationWarning)
        z_grid = mesh.nodes_x
        n_sig = sig.size
        sigma = np.repeat(sig, 2)
        z = []
        for i in range(n_sig):
            z.append(np.r_[z_grid[i], z_grid[i + 1]])
        z = np.hstack(z)
        if xlim == None:
            sig_min = sig[~np.isnan(sig)].min() * 0.5
            sig_max = sig[~np.isnan(sig)].max() * 2
        else:
            sig_min, sig_max = xlim

        if xscale == "linear" and sig.min() == 0.0:
            if xlim == None:
                sig_min = -sig[~np.isnan(sig)].max() * 0.5
                sig_max = sig[~np.isnan(sig)].max() * 2

        if ax == None:
            plt.xscale(xscale)
            plt.xlim(sig_min, sig_max)
            plt.ylim(z.min(), z.max())
            plt.xlabel("Conductivity (S/m)")
            plt.ylabel("Depth (m)")
            plt.ylabel("Depth (m)")
            if showlayers == True:
                for locz in z_grid:
                    plt.plot(
                        np.linspace(sig_min, sig_max, 100),
                        np.ones(100) * locz,
                        "b--",
                        lw=0.5,
                    )
            return plt.plot(sigma, z, "k-", **kwargs)

        else:
            ax.set_xscale(xscale)
            ax.set_xlim(sig_min, sig_max)
            ax.set_ylim(z.min(), z.max())
            ax.set_xlabel("Conductivity (S/m)")
            ax.set_ylabel("Depth (m)")
            if showlayers == True:
                for locz in z_grid:
                    ax.plot(
                        np.linspace(sig_min, sig_max, 100),
                        np.ones(100) * locz,
                        "b--",
                        lw=0.5,
                    )
            return ax.plot(sigma, z, "k-", **kwargs)


def set_mesh_1d(hz):
    return TensorMesh([hz], x0=[0])
