from __future__ import print_function
from __future__ import absolute_import
from __future__ import unicode_literals

import numpy as np
from SimPEG import EM
import matplotlib.pyplot as plt
import matplotlib.ticker as ticker
import matplotlib
import matplotlib.gridspec as gridspec

from ipywidgets import widgets
from scipy.constants import mu_0, epsilon_0

from .DipoleWidgetTD import DipoleWidgetTD, linefun, DisPosNegvalues
from .view import DataView
from .VolumeWidget import polyplane
from .TDEMPlanewave import (
    e_field_from_sheet_current,
    h_field_from_sheet_current,
    j_field_from_sheet_current,
)

matplotlib.rcParams["font.size"] = 12


def PlaneEHfield(z, t=0.0, sig=1.0, mu=mu_0, epsilon=epsilon_0, E0=1.0):
    """
        Plane wave propagating downward (negative z (depth))
    """
    bunja = -E0 * (mu * sig) ** 0.5 * z * np.exp(-(mu * sig * z ** 2) / (4 * t))
    bunmo = 2 * np.pi ** 0.5 * t ** 1.5
    Ex = bunja / bunmo
    Hy = E0 * np.sqrt(sig / (np.pi * mu * t)) * np.exp(-(mu * sig * z ** 2) / (4 * t))
    return Ex, Hy


class PlanewaveWidget(DipoleWidgetTD):
    """PlanewaveWidget"""

    x = None
    y = None
    z = None
    func = None

    # Fixed spatial range in 3D
    xmin, xmax = -500.0, 500.0
    ymin, ymax = -500.0, 500.0
    zmin, zmax = -1000.0, 0.0

    def __init__(self):
        self.dataview = DataView()

    def SetDataview(
        self, srcLoc, sig, t, orientation, normal, functype, na=100, nb=100, loc=0.0
    ):

        self.srcLoc = srcLoc
        self.sig = sig
        self.t = t
        self.normal = normal
        self.SetGrid(normal, loc, na, nb)
        self.functype = functype
        self.dataview.set_xyz(
            self.x, self.y, self.z, normal=normal
        )  # set plane and locations ...

        if self.functype == "E_from_SheetCurrent":
            self.func = e_field_from_sheet_current
        elif self.functype == "H_from_SheetCurrent":
            self.func = h_field_from_sheet_current
        elif self.functype == "J_from_SheetCurrent":
            self.func = j_field_from_sheet_current
        else:
            raise NotImplementedError()

        self.dataview.eval_2D_TD(srcLoc, sig, t, orientation, self.func)  # evaluate

    def Planewave2Dviz(
        self,
        x1,
        y1,
        x2,
        y2,
        npts2D,
        npts,
        sig,
        t,
        srcLoc=0.0,
        orientation="x",
        view="x",
        normal="Z",
        functype="E_from_ED",
        loc=0.0,
        scale="log",
        dx=50.0,
    ):
        nx, ny = npts2D, npts2D
        x, y = linefun(x1, x2, y1, y2, npts)
        if scale == "log":
            logamp = True
        elif scale == "linear":
            logamp = False
        else:
            raise NotImplementedError()
        self.SetDataview(
            srcLoc, sig, t, orientation, normal, functype, na=nx, nb=ny, loc=loc
        )
        plot1D = True
        if normal == "X" or normal == "x":
            xyz_line = np.c_[np.ones_like(x) * self.x, x, y]
            self.dataview.xyz_line = xyz_line
        if normal == "Y" or normal == "y":
            xyz_line = np.c_[x, np.ones_like(x) * self.y, y]
            self.dataview.xyz_line = xyz_line
        if normal == "Z" or normal == "z":
            xyz_line = np.c_[x, y, np.ones_like(x) * self.z]
            self.dataview.xyz_line = xyz_line

        plt.figure(figsize=(18 * 1.5, 3.4 * 1.5))
        gs1 = gridspec.GridSpec(2, 7)
        gs1.update(left=0.05, right=0.48, wspace=0.05)
        ax1 = plt.subplot(gs1[:2, :3])
        ax1.axis("equal")

        ax1, dat1 = self.dataview.plot2D_TD(
            ax=ax1, view=view, colorbar=False, logamp=logamp
        )
        vmin, vmax = dat1.cvalues.min(), dat1.cvalues.max()
        if scale == "log":
            cb = plt.colorbar(
                dat1, ax=ax1, ticks=np.linspace(vmin, vmax, 5), format="$10^{%.1f}$"
            )
        elif scale == "linear":
            cb = plt.colorbar(
                dat1, ax=ax1, ticks=np.linspace(vmin, vmax, 5), format="%.1e"
            )

        tempstr = functype.split("_")

        title = tempstr[0] + view

        if tempstr[0] == "E":
            unit = " (V/m)"
            fieldname = "Electric field"
        elif tempstr[0] == "H":
            unit = " (A/m)"
            fieldname = "Magnetic field"
        else:
            raise NotImplementedError()

        label = fieldname + unit
        cb.set_label(label)
        ax1.set_title(title)

        if plot1D:
            ax1.plot(x, y, "r.", ms=4)
            ax2 = plt.subplot(gs1[:, 4:6])
            val_line_x, val_line_y, val_line_z = self.dataview.eval_TD(
                xyz_line, srcLoc, np.r_[sig], np.r_[t], orientation, self.func
            )

            if view == "X" or view == "x":
                val_line = val_line_x
            elif view == "Y" or view == "y":
                val_line = val_line_y
            elif view == "Z" or view == "z":
                val_line = val_line_z

            distance = xyz_line[:, 2]

            if scale == "log":
                temp = val_line.copy() * np.nan
                temp[val_line > 0.0] = val_line[val_line > 0.0]
                ax2.plot(temp, distance, "k.-")
                temp = val_line.copy() * np.nan
                temp[val_line < 0.0] = -val_line[val_line < 0.0]
                ax2.plot(temp, distance, "k.--")
                ax2.set_xlim(abs(val_line).min(), abs(val_line).max())
                ax2.set_xscale(scale)

            elif scale == "linear":
                ax2.plot(val_line, distance, "k.-")
                ax2.set_xlim(val_line.min(), val_line.max())
                ax2.set_xscale(scale)
                xticks = np.linspace(-abs(val_line).max(), abs(val_line).max(), 3)
                plt.plot(
                    np.r_[0.0, 0.0], np.r_[distance.min(), distance.max()], "k-", lw=2
                )
                ax2.xaxis.set_ticks(xticks)
                ax2.xaxis.set_major_formatter(ticker.FormatStrFormatter("%.0e"))
                ax2.set_xlim(-abs(val_line).max(), abs(val_line).max())

            ax2.set_ylim(distance.min(), distance.max())
            ax2.set_ylabel("Profile (m)")

            if tempstr[0] == "E":
                label = "(" + tempstr[0] + view + ")-field (V/m) "
            elif tempstr[0] == "H":
                label = "(" + tempstr[0] + view + ")-field (A/m) "
            elif tempstr[0] == "J":
                label = "(" + tempstr[0] + view + ")-field (A/m$^2$) "
            else:
                raise NotImplementedError()

            ax2.set_title("EM data")
            ax2.set_xlabel(label)
            ax2.grid(True)
        plt.show()
        pass

    def InteractivePlaneWave(
        self,
        nRx=100,
        npts2D=50,
        scale="log",
        offset_plane=50.0,
        X1=-500,
        X2=500,
        Y1=-500,
        Y2=500,
        Z1=-1000,
        Z2=0,
    ):

        # x1, x2, y1, y2 = offset_rx, offset_rx, Z1, Z2
        self.xmin, self.xmax = X1, X2
        self.ymin, self.ymax = Y1, Y2
        self.zmin, self.zmax = Z1, Z2

        def foo(Field, Time, Sigma, Scale):
            t = np.r_[Time]
            sig = np.r_[Sigma]

            if Field == "Ex":
                normal = "Y"
                self.offset_rx = 0.0
                Field = "E_from_SheetCurrent"
                Component = "x"

            elif Field == "Hy":
                normal = "X"
                self.offset_rx = 0.0
                Field = "H_from_SheetCurrent"
                Component = "y"

            x1, x2, y1, y2 = self.offset_rx, self.offset_rx, Z1, Z2

            return self.Planewave2Dviz(
                x1,
                y1,
                x2,
                y2,
                npts2D,
                nRx,
                sig,
                t,
                srcLoc=0.0,
                orientation="X",
                view=Component,
                normal=normal,
                functype=Field,
                scale=Scale,
            )

        out = widgets.interactive(
            foo,
            Field=widgets.ToggleButtons(options=["Ex", "Hy"]),
            Time=widgets.FloatSlider(
                min=0.01, max=1.0, step=0.01, value=0.0, description="$t$ (s)"
            ),
            Sigma=widgets.FloatText(value=1, continuous_update=False),
            Scale=widgets.ToggleButtons(options=["log", "linear"], value="linear"),
        )
        return out


def InteractivePlaneProfile():
    nRx = 100

    def foo(Field, Sigma, Scale, Time):

        plt.figure(figsize=(8, 4))
        ax1 = plt.subplot(111)

        r = np.linspace(-1000.0, 0.0, nRx)
        val_ex, val_hy = PlaneEHfield(r, t=Time, sig=Sigma)

        if Field == "Ex":
            val = val_ex.flatten()
            label = "Ex-field (V/m)"

        elif Field == "Hy":
            val = val_hy.flatten()
            label = "Hy-field (A/m)"

        if Scale == "log":
            val_p, val_n = DisPosNegvalues(val)
            ax1.plot(r, val_p, "k-", lw=2)
            ax1.plot(r, val_n, "k--", lw=2)
            ax1.set_yscale(Scale)

        elif Scale == "linear":
            ax1.plot(r, val, "k-", lw=2)
            ax1.set_yscale(Scale)
            y = ax1.yaxis.get_majorticklocs()
            yticksa = np.linspace(y.min(), y.max(), 3)
            ax1.yaxis.set_ticks(yticksa)
            ax1.yaxis.set_major_formatter(ticker.FormatStrFormatter("%.1e"))

        ax1.set_xlim(0, -1000)
        ax1.set_ylabel(label, color="k")
        ax1.set_xlabel("Z (m)")

        ax1.grid(True)
        plt.show()

    Q2 = widgets.interactive(
        foo,
        Field=widgets.ToggleButtons(options=["Ex", "Hy"], value="Ex"),
        Sigma=widgets.FloatText(
            value=1, continuous_update=False, description="$\sigma$ (S/m)"
        ),
        Scale=widgets.ToggleButtons(options=["log", "linear"], value="linear"),
        Time=widgets.FloatSlider(
            min=0.01, max=1.0, step=0.01, value=0.0, description="$t$ (s)"
        ),
    )
    return Q2
