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

import numpy as np
from SimPEG.electromagnetics.utils import waveform_utils
import matplotlib.pyplot as plt
import matplotlib.ticker as ticker
import matplotlib
import matplotlib.gridspec as gridspec
from ipywidgets import widgets, interactive, IntSlider
from scipy.constants import mu_0, epsilon_0

from .DipoleWidgetFD import DipoleWidgetFD, linefun, DisPosNegvalues
from .VolumeWidget import polyplane
from .FDEMPlanewave import (
    e_field_from_sheet_current,
    b_field_from_sheet_current,
    h_field_from_sheet_current,
    j_field_from_sheet_current,
)
from .view import DataView

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


def PlaneEHfield(z, t=0.0, f=1.0, sig=1.0, mu=mu_0, epsilon=epsilon_0, E0=1.0):
    """
        Plane wave propagating downward (negative z (depth))
    """
    k = waveform_utils.k(f, sig, mu=mu, eps=epsilon)
    omega = waveform_utils.omega(f)
    Ex = E0 * np.exp(1j * (k * z + omega * t))
    Z = omega * mu / k
    Hy = -E0 / Z * np.exp(1j * (k * z + omega * t))
    return Ex, Hy


class PolarEllipse(object):
    def Planewave3D(self, itime):
        fig = plt.figure(figsize=(12 * 1.2, 5 * 1.2))
        ax = fig.add_subplot(111, projection="3d")
        t = np.linspace(0, 0.02, 1000)
        X1, X2 = t.min(), t.max()
        Y1, Y2 = -1.2, 1.2
        Z1, Z2 = -1.2, 1.2

        Ex, Hy = PlaneEHfield(-10.0, t=t, f=100.0, sig=1e-2)
        ax.plot(t, Ex.real / Ex.real.max(), np.zeros_like(t), "b", lw=1)
        ax.plot(np.r_[t[itime]], (Ex[itime]).real / Ex.real.max(), 0.0, "bo", ms=5)
        ax.plot(
            t[itime] * np.ones(2),
            np.r_[0.0, (Ex[itime]).real / Ex.real.max()],
            np.ones(2) * (Hy[itime]).real / Hy.real.max(),
            "r:",
        )
        ax.plot(t, np.zeros_like(t), Hy.real / Hy.real.max(), "r", lw=1)
        ax.plot(np.r_[t[itime]], 0.0, (Hy[itime]).real / Hy.real.max(), "ro", ms=5)
        ax.plot(
            t[itime] * np.ones(2),
            np.ones(2) * (Ex[itime]).real / Ex.real.max(),
            np.r_[0.0, (Hy[itime]).real / Hy.real.max()],
            "b:",
        )

        ax.plot(
            np.ones_like(t) * t[itime],
            Ex.real / Ex.real.max(),
            Hy.real / Hy.real.max(),
            "k-",
            lw=0.5,
        )
        ax.plot(
            np.ones(1) * t[itime],
            Ex[itime].real / Ex.real.max(),
            Hy[itime].real / Hy.real.max(),
            "ko",
            ms=5,
        )
        ax.plot(t, np.zeros_like(t), np.zeros_like(t), "k--")
        ax.plot(np.ones(2) * t[itime], np.r_[Y1, Y2], np.zeros(2), "k--")
        ax.plot(np.ones(2) * t[itime], np.zeros(2), np.r_[Z1, Z2], "k--")
        x = np.r_[X1, X2, X2, X1, X1]
        y = np.zeros(5)
        z = np.r_[Z1, Z1, Z2, Z2, Z1]
        verts = zip(x, y, z)
        polya = polyplane(verts, color="red", alpha=0.1)
        ax.plot(x, y, z, "r-", lw=1, alpha=0.2)
        x = np.r_[X1, X2, X2, X1, X1]
        y = np.r_[Y1, Y1, Y2, Y2, Y1]
        z = np.zeros(5)
        verts = zip(x, y, z)
        polyb = polyplane(verts, color="blue", alpha=0.1)
        ax.plot(x, y, z, "b-", lw=1, alpha=0.2)
        x = np.ones(5) * t[itime]
        y = np.r_[Y1, Y2, Y2, Y1, Y1]
        z = np.r_[Z1, Z1, Z2, Z2, Z1]
        verts = zip(x, y, z)
        polyc = polyplane(verts, color="grey", alpha=0.1)
        ax.plot(x, y, z, "k-", lw=1, alpha=0.2)

        ax.add_collection3d(polya)
        ax.add_collection3d(polyb)
        ax.add_collection3d(polyc)

        ax.set_zlim(1.2, -1.2)
        ax.set_ylim(-1.2, 1.2)
        ax.set_xlim(t.min(), t.max())
        ax.set_xlabel("Time (s)")
        ax.set_ylabel("Ex (V/m)")
        ax.set_zlabel("Hy (A/m)")
        elev = 45
        azim = 290
        ax.view_init(elev, azim)
        plt.show()
        pass

    def Interactive(self):
        out = interactive(self.Planewave3D, itime=IntSlider(min=0, max=999, step=10))
        return out


class PlanewaveWidget(DipoleWidgetFD):
    """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,
        f,
        orientation,
        normal,
        functype,
        na=100,
        nb=100,
        loc=0.0,
        t=0.0,
    ):

        self.srcLoc = srcLoc
        self.sig = sig
        self.f = f
        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(srcLoc, sig, f, orientation, self.func, t=t)  # evaluate

    def Planewave2Dviz(
        self,
        x1,
        y1,
        x2,
        y2,
        npts2D,
        npts,
        sig,
        f,
        srcLoc=0.0,
        orientation="x",
        component="real",
        view="x",
        normal="Z",
        functype="E_from_ED",
        loc=0.0,
        scale="log",
        dx=50.0,
        t=0.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, f, orientation, normal, functype, na=nx, nb=ny, loc=loc, t=t
        )
        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_FD(
            ax=ax1, component=component, 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("_")

        tname_end = ")"

        if component == "real":
            tname = "Re("
        elif component == "imag":
            tname = "Im("
        elif component == "amplitude":
            tname = "|"
            tname_end = "|"
        elif component == "phase":
            tname = "Phase("

        title = tname + tempstr[0] + view + tname_end

        if tempstr[0] == "E":
            unit = " (V/m)"
            fieldname = "Electric field"
        elif tempstr[0] == "H":
            unit = " (A/m)"
            fieldname = "Magnetic field"
        elif tempstr[0] == "J":
            unit = " (A/m$^2$) "
            fieldname = "Current density"
        else:
            raise NotImplementedError()
        if component == "phase":
            unit = " (rad)"
        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(
                xyz_line, srcLoc, np.r_[sig], np.r_[f], orientation, self.func, t=t
            )

            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
            elif view == "vec" or "amp":

                def vecamp(a, b, c):
                    return np.sqrt((a) ** 2 + (b) ** 2 + (c) ** 2)

                if component == "real":
                    val_line = vecamp(val_line_x.real, val_line_y.real, val_line_z.real)
                elif component == "imag":
                    val_line = vecamp(val_line_x.imag, val_line_y.imag, val_line_z.imag)
                elif component == "amplitude":
                    val_line = vecamp(abs(val_line_x), abs(val_line_y), abs(val_line_z))
                elif component == "phase":
                    val_line = vecamp(
                        np.angle(val_line_x), np.angle(val_line_y), np.angle(val_line_z)
                    )

            distance = xyz_line[:, 2]

            if component == "real":
                val_line = val_line.real
            elif component == "imag":
                val_line = val_line.imag
            elif component == "amplitude":
                val_line = abs(val_line)
            elif component == "phase":
                val_line = np.angle(val_line)

            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":
                if view == "vec" or view == "amp":
                    label = "|" + tempstr[0] + "|-field (V/m) "
                else:
                    label = tname + tempstr[0] + view + ")-field (V/m) "
            elif tempstr[0] == "H":
                if view == "vec" or view == "amp":
                    label = "|" + tempstr[0] + "|-field field (A/m) "
                else:
                    label = tname + tempstr[0] + view + ")-field (A/m) "
            elif tempstr[0] == "J":
                if view == "vec" or view == "amp":
                    label = "|" + tempstr[0] + "|-field field (A/m$^2$) "
                else:
                    label = tname + tempstr[0] + view + ")-field (A/m$^2$) "
            else:
                raise NotImplementedError()

            if component == "phase":
                label = tname + tempstr[0] + view + ")-field (rad) "

            ax2.set_title("EM data")
            ax2.set_xlabel(label)

            # ax2.text(distance.min(), val_line.max(), 'A', fontsize = 16)
            # ax2.text(distance.max()*0.97, val_line.max(), 'B', fontsize = 16)
            # ax2.legend((component, ), bbox_to_anchor=(0.5, -0.3))
            ax2.grid(True)
        # plt.tight_layout()
        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, ComplexNumber, Frequency, Sigma, Scale, Time):
            f = np.r_[Frequency]
            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

            if ComplexNumber == "Re":
                ComplexNumber = "real"
            elif ComplexNumber == "Im":
                ComplexNumber = "imag"
            elif ComplexNumber == "Amp":
                ComplexNumber = "amplitude"
            elif ComplexNumber == "Phase":
                ComplexNumber = "phase"

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

        out = widgets.interactive(
            foo,
            Field=widgets.ToggleButtons(options=["Ex", "Hy"]),
            ComplexNumber=widgets.ToggleButtons(options=["Re", "Im", "Amp", "Phase"]),
            Frequency=widgets.FloatText(value=10.0, continuous_update=False),
            Sigma=widgets.FloatText(value=1, continuous_update=False),
            Scale=widgets.ToggleButtons(options=["log", "linear"], value="linear"),
            Time=widgets.FloatSlider(min=0, max=0.2, step=0.01, value=0.0),
        )
        return out


def InteractivePlaneProfile():
    # srcLoc = 0.0
    # orientation = "X"
    nRx = 100

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

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

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

        if Field == "Ex":
            valr = val_ex.real.flatten()
            vali = val_ex.imag.flatten()
            labelr = "Re (Ex)-field (V/m)"
            labeli = "Im (Ex)-field (V/m)"
        elif Field == "Hy":
            valr = val_hy.real.flatten()
            vali = val_hy.imag.flatten()
            labelr = "Re (Hy)-field (A/m)"
            labeli = "Im (Hy)-field (A/m)"

        elif Field == "Impedance":
            imp = -val_ex / val_hy
            valr = imp.real.flatten()
            vali = imp.imag.flatten()
            labelr = "Re (Z) (Ohm)"
            labeli = "Im (Z) (Ohm)"

        elif Field == "rhophi":
            imp = -val_ex / val_hy
            valr = abs(imp) ** 2 / (2 * np.pi * Frequency * mu_0)
            vali = np.angle(imp, deg=True)
            labelr = "Apparent resistivity (Ohm-m)"
            labeli = "Phase of Impedance (degree)"

        if Scale == "log":
            valr_p, valr_n = DisPosNegvalues(valr)
            vali_p, vali_n = DisPosNegvalues(vali)
            if Field == "rhophi":
                ax1.plot(r, valr, "k.")
            else:
                ax1.plot(r, valr_p, "k-", lw=2)
            ax1.plot(r, valr_n, "k--", lw=2)
            ax2.plot(r, vali_p, "r-", lw=2)
            ax2.plot(r, vali_n, "r--", lw=2)
            ax1.set_yscale(Scale)
            ax2.set_yscale(Scale)
            if Fixed:
                vmin1, vmax1 = ax1.get_ylim()
                vmin2, vmax2 = ax2.get_ylim()
                vmin = min(vmin1, vmin2)
                vmax = max(vmax1, vmax2)
                ax1.set_ylim(vmin, vmax)
                ax2.set_ylim(vmin, vmax)

        elif Scale == "linear":
            if Field == "rhophi":
                ax1.plot(r, valr, "k.")
            else:
                ax1.plot(r, valr, "k-", lw=2)
            ax2.plot(r, vali, "r-", lw=2)
            ax1.set_yscale(Scale)
            ax2.set_yscale(Scale)
            y = ax1.yaxis.get_majorticklocs()
            yticksa = np.linspace(y.min(), y.max(), 3)
            ax1.yaxis.set_ticks(yticksa)

            if Fixed and Field != "Impedance":
                vmax = np.r_[abs(valr), abs(vali)].max()
                vmin = -vmax
                ax1.set_ylim(vmin, vmax)
                ax2.set_ylim(vmin, vmax)
                # y = ax2.yaxis.get_majorticklocs()
                yticks = np.linspace(vmin, vmax, 3)
                ax1.yaxis.set_ticks(yticks)
                ax2.yaxis.set_ticks(yticks)

            ax1.yaxis.set_major_formatter(ticker.FormatStrFormatter("%.1e"))
            ax2.yaxis.set_major_formatter(ticker.FormatStrFormatter("%.1e"))

        ax1.set_xlim(0, -1000)
        ax2.set_ylabel(labeli, color="r")
        ax1.set_ylabel(labelr, color="k")
        ax1.set_xlabel("Z (m)")

        for tl in ax2.get_yticklabels():
            tl.set_color("r")
        ax1.grid(True)
        plt.show()

    Q2 = widgets.interactive(
        foo,
        Field=widgets.ToggleButtons(
            options=["Ex", "Hy", "Impedance", "rhophi"], value="Ex"
        ),
        Sigma=widgets.FloatText(
            value=1, continuous_update=False, description="$\sigma$ (S/m)"
        ),
        Scale=widgets.ToggleButtons(options=["log", "linear"], value="linear"),
        Fixed=widgets.widget_bool.Checkbox(value=False),
        Frequency=widgets.FloatText(
            value=10.0, continuous_update=False, description="$f$ (Hz)"
        ),
        Time=widgets.FloatSlider(
            min=0, max=0.2, step=0.005, continuous_update=False, description="t (s)"
        ),
    )
    return Q2
