# -*- coding: utf-8 -*-


import numpy as np
import numpy.linalg as la
from warnings import warn
from typing import List, Union


class S2T:

    def __init__(self, xmin, xmax, dx, ymin, ymax, dy, epsi_bg, mu_bg, x=None, y=None):

        """
        Set up the structure to be solved in 2D FDTD.

        Parameters
        ----------
        xmin, xmax, dx, ymin, ymax, dy      :   float
        epsi_bg, mu_bg                      :   Union[float, np.ndarray]
        x, y                                :   ndarray
                                                1d ndarray, If supplied, override xmin etc.
        """

        if type(epsi_bg) is float:
            self.epsi_bg = np.eye(3, 3) * epsi_bg
        elif np.array(epsi_bg).ndim == 1:
            self.epsi_bg = np.diag(epsi_bg)
        else:
            self.epsi_bg = epsi_bg
        if type(mu_bg) is float:
            self.mu_bg = np.eye(3, 3) * mu_bg
        elif np.array(mu_bg).ndim == 1:
            self.mu_bg = np.diag(mu_bg)
        else:
            self.mu_bg = mu_bg
        self.epsi_bg_inv = la.inv(self.epsi_bg)
        self.mu_bg_inv = la.inv(self.mu_bg)

        # parameters of the solving space
        if x is not None:
            dx = x[1] - x[0]
            xmin = x[0]
            xmax = x[-1] + dx
            Nx = x.size
            x_n = np.arange(Nx, dtype=int)
        else:
            Nx = int((xmax + dx / 1.e4 - xmin) / dx)
            x_n = np.arange(Nx, dtype=int)
            x = x_n * dx + xmin
        self.xmin = xmin
        self.xmax = xmax
        self.dx = dx
        self.xres = dx
        self.Nx = Nx
        self.x_n = x_n
        self.x = x
        #
        if y is not None:
            dy = y[1] - y[0]
            ymin = y[0]
            ymax = y[-1] + dy
            Ny = y.size
            y_n = np.arange(Nx, dtype=int)
        else:
            Ny = int((ymax + dx / 1.e4 - ymin) / dy)
            y_n = np.arange(Ny, dtype=int)
            y = y_n * dy + ymin
        self.ymin = ymin
        self.ymax = ymax
        self.dy = dy
        self.yres = dy
        self.Ny = Ny
        self.y_n = y_n
        self.y = y
        #
        self.xx_n, self.yy_n = np.meshgrid(self.x_n, self.y_n)
        self.xx, self.yy = np.meshgrid(self.x, self.y)

        # list of box objects
        self.bxs: List[Bx2] = []
        self.bxs_sp: List[Bx2] = []
        self.bxs_Lz: List[Bx2] = []
        self.bxs_dmLz: List[Bx2] = []
        self.bxs_Dr: List[Bx2] = []
        self.bxs_dmri: List[Bx2] = []
        self.bxs_lg: List[Bx2] = []
        self.bxs_dmlg: List[Bx2] = []

    def a_b(self,
            shp='rct',
            xmin_b=None, xmax_b=None, ymin_b=None, ymax_b=None,
            xc_b=None, yc_b=None,
            r_b=None,
            a_1=None, a_2=None,
            w_b=None,
            knd='sp',
            epsi_b=None, mu_b=None,
            epsi_infty_b=None, mu_infty_b=None,
            omgs_rsn=None, dlts_epsi=None, Gms=None,
            sgm=None,
            m_a=None, m_o=None, m_q=None, m_p=0.):

        """
        Add boxes to the structure.

        Parameters
        ----------
        shp                             :   str
                                            ['rct'|'ccl'|'rng'|'wdg']
        xmin_b, xmax_b, ymin_b, ymax_b  :   float
                                            four corners of the box.
        xc_b, yc_b                      :   float
                                            x and y center of box
        r_b                             :   float
                                            radius of circular box, outer radius for ring and wedge box
        a_1, a_2                        :   float
                                            start and end angle of wedge, in degrees, between 0 to 360
        w_b                             :   float
                                            width of ring or wedge box
        knd                             :   str
                                            ['sp'|'Lz'|'dmLz'|'Dr'|'lg'|'dmri'|'dmlg']
        epsi_b, mu_b                    :   float or Tuple[float, float, float] or 3x3 ndarray
                                            The permittivity and permeability.
        epsi_infty_b, mu_infty_b        :   float or Tuple[float, float, float] or 3x3 ndarray
                                            The high-frequency permittivity and permeability.
                                            If not supplied, fall back to epsi_b and mu_b.
        omgs_rsn                        :   list[float]
                                            resonance frequencies
        dlts_epsi                       :   list[float]
                                            permittivity change at each resonance
        Gms                             :   list[float]
                                            Broadening of each resonance
        sgm                             :   float
                                            loss (positive number) or gain (negative) in the material
        m_a                             :   float or list[float]
                                            modulation amplitude
        m_o                             :   float or list[float]
                                            modulation frequency
        m_q                             :   Tuple[float, float] or list[Tuple[float, float]]
                                            (qx, qy), wave vector of modulation.
        m_p                             :   float or list[float]
                                            modulation phase
                                            By default, cos modulation.

        Returns
        -------

        """

        gmt = Gmt2(self.xx, self.yy,
                   shp=shp,
                   xmin_rec_b=xmin_b, xmax_rec_b=xmax_b, ymin_rec_b=ymin_b, ymax_rec_b=ymax_b,
                   xc_b=xc_b, yc_b=yc_b,
                   r_b=r_b,
                   a_1=a_1, a_2=a_2,
                   w_b=w_b)

        mtr = Mtr(knd=knd,
                  epsi=epsi_b, mu=mu_b, epsi_infty=epsi_infty_b, mu_infty=mu_infty_b,
                  omgs_rsn=omgs_rsn, dlts_epsi=dlts_epsi, Gms=Gms,
                  sgm=sgm,
                  m_a=m_a, m_o=m_o, m_q=m_q, m_p=m_p)

        bx = Bx2(gmt, mtr)

        if knd == 'sp':
            self.bxs_sp.append(bx)
        elif knd == 'Lz':
            self.bxs_Lz.append(bx)
        elif knd == 'dmLz':
            self.bxs_dmLz.append(bx)
        elif knd == 'Dr':
            self.bxs_Dr.append(bx)
        elif knd == 'dmri':
            self.bxs_dmri.append(bx)
        elif knd == 'lg':
            self.bxs_lg.append(bx)
        elif knd == 'dmlg':
            self.bxs_dmlg.append(bx)
        else:
            print('Material not recognized.')

        if bx is not None:
            self.bxs.append(bx)


class Bx2:
    def __init__(self, gmt, mtr):
        """
        A box in 2D.

        Parameters
        ----------
        gmt        :   Gmt2
        mtr        :   Mtr
        """

        self.gmt = gmt
        self.mtr = mtr


class MtrBsc(object):
    def __init__(self, epsi=None, mu=None, **kwargs):
        """
        Basic material.

        Parameters
        ----------
        epsi, mu    :   float | (float) | numpy.ndarray
                        permittivity and permeability. if tuple, it is length of 3. if ndarray, it has shape (3, 3)
        kwargs      :
                        other kwargs not related to materials
        """
        if epsi is None:
            epsi = 1.
        if mu is None:
            mu = 1.

        if not hasattr(epsi, '__len__'):
            epsi = np.eye(3, 3) * epsi
        elif np.array(epsi).ndim == 1:
            epsi = np.diag(epsi)
        if not hasattr(mu, '__len__'):
            mu = np.eye(3, 3) * mu
        elif np.array(mu).ndim == 1:
            mu = np.diag(mu)
        self.epsi = epsi
        self.mu = mu


class Mtr(MtrBsc):
    def __init__(self,
                 knd='sp',
                 epsi_infty=1., mu_infty=1., omgs_rsn=None, dlts_epsi=None, Gms=None,
                 sgm=None,
                 m_a=None, m_o=None, m_q=None, m_p=None,
                 **kw):
        """
        Material type and parameters.

        Parameters
        ----------
        knd                             :   str
                                            Options are:
                                            'sp' (simple),
                                            'Lz' (Lorentz),
                                            'dmLz' (dynamic modulation of delta epsi),
                                            'Dr' (Drude),
                                            'lg' (loss and gain),
                                            'dmri' (dynamic modulation of real part of index),
                                            'dmlg' (dynamic modulation of loss and gain).
        epsi_infty_b                    :   float | (float) | numpy.ndarray
        mu_infty_b                      :   float | (float) | numpy.ndarray
                                            if tuple it has len 3, if ndarray it has shape (3, 3).

                                            The high-frequency permittivity and permeability for use in Lorentz and Drude.

                                            If not supplied, fall back to `epsi_b` and `mu_b`.

        omgs_rsn                        :   list[float]
                                            resonance frequencies, for use in Lorentz or Drude
        dlts_epsi                       :   list[float]
                                            permittivity change at each resonance, for use in Lorentz
        Gms                             :   list[float]
                                            Broadening of each resonance, for use in Lorentz or Drude
        sgm                             :   float
                                            loss (positive number) or gain (negative) in the material
        m_a                             :   float or list[float]
                                            modulation amplitude
        m_o                             :   float or list[float]
                                            modulation frequency
        m_q                             :   Tuple[float, float] or list[Tuple[float, float]]
                                            (qx, qy), wave vector of modulation.
        m_p                             :   float or list[float]
                                            modulation phase
                                            By default, cos modulation.

        Keyword Arguments
        -----------------
        kw                      :   keyword arguments to pass to MtrBsc
        """

        if (knd == 'Lz') or (knd == 'dmLz') or (knd == 'Dr'):
            if epsi_infty is None:
                epsi_infty = kw.pop('epsi')
            if mu_infty is None:
                mu_infty = kw.pop('mu')
            super(Mtr, self).__init__(epsi=epsi_infty, mu=mu_infty)
            self.epsi_infty = self.epsi
            self.mu_infty = self.mu
        else:
            super(Mtr, self).__init__(**kw)

        self.knd = knd

        if omgs_rsn is None:
            omgs_rsn = []
        if dlts_epsi is None:
            dlts_epsi = []
        if Gms is None:
            Gms = []

        self.omgs_rsn = omgs_rsn
        self.dlts_epsi = dlts_epsi
        self.Gms = Gms

        self.sgm = sgm

        self.m_a = m_a
        self.m_o = m_o
        self.m_q = m_q
        self.m_p = m_p


class Gmt2:
    def __init__(self, xx, yy,
                 shp='rct',
                 xmin_rec_b=None, xmax_rec_b=None, ymin_rec_b=None, ymax_rec_b=None,
                 xc_b=None, yc_b=None,
                 r_b=None,
                 w_b=None,
                 a_1=None, a_2=None,
                 **kwargs
                 ):

        """
        Geometrical information of a box.

        Parameters
        ----------
        xx, yy          :   ndarray
        shp             :   str
                            'rct' (rectangular),
                            'ccl' (circular),
                            'rng' (ring),
                            'wdg' (wedge)
        xmin_rec_b
        xmax_rec_b
        ymin_rec_b
        ymax_rec_b      :   float
                            four corners of a rectangular box.
        xc_b, yc_b      :   float
                            x and y center of box
        r_b             :   float
                            radius of circular box, the outer radius for ring and wedge box
        a_1, a_2        :   float
                            the start and the end angle of an wedge, in degrees, between 0 to 360
        w_b             :   float
                            width of a ring or wedge box
        kwargs          :
                            other kwargs supplied but not related to geometry
        """

        # self.shp = shp
        if shp == 'rct':
            self.pch = {'shp': 'rct',
                        'xy': (xmin_rec_b, ymin_rec_b),
                        'width': xmax_rec_b - xmin_rec_b,
                        'height': ymax_rec_b - ymin_rec_b
                        }
            # indexing in box
            self.iib = (xx < xmax_rec_b) * (xx >= xmin_rec_b) * (yy < ymax_rec_b) * (yy >= ymin_rec_b)
            self.iibV = self.iib.ravel()

        elif shp == 'ccl':
            self.pch = {'shp': 'ccl',
                        'center': (xc_b, yc_b),
                        'radius': r_b,
                        }
            # indexing in box
            d = np.sqrt(((xx - xc_b) ** 2 + (yy - yc_b) ** 2))
            self.iib = d < r_b
            self.iibV = self.iib.ravel()

        elif shp == 'rng':
            self.pch = {'shp': 'rng',
                        'center': (xc_b, yc_b),
                        'radius': r_b,
                        'width': w_b,
                        }
            # indexing in box
            d = np.sqrt(((xx - xc_b) ** 2 + (yy - yc_b) ** 2))
            self.iib = (d < r_b) * (d >= r_b - w_b)
            self.iibV = self.iib.ravel()

        elif shp == 'wdg':
            self.pch = {'shp': 'wdg',
                        'center': (xc_b, yc_b),
                        'radius': r_b,
                        'angles': (a_1, a_2),
                        'width': w_b,
                        }
            # indexing in box
            d = np.sqrt(((xx - xc_b) ** 2 + (yy - yc_b) ** 2))
            if w_b is None:
                i_i_r = d < r_b
            else:
                i_i_r = (d < r_b) * (d >= r_b - w_b)
            z_ = (xx - xc_b) + (yy - yc_b) * 1j
            a_ = np.angle(z_, deg=True)
            a1 = (a_1 + 180) % 360 - 180
            # regularize angles
            if a_2 != 180.:
                a2 = (a_2 + 180) % 360 - 180
            else:
                a2 = 180.
            if a2 > a1:
                i_i_a = (a_ >= a1) * (a_ < a2)
            else:
                i_i_a = np.logical_not(((a_ >= a1) * (a_ < a2)))
            self.iib = i_i_r * i_i_a
            self.iibV = self.iib.ravel()
