import numpy as np
import pyproj
import scipy.ndimage
from scipy import interpolate
from netCDF4 import Dataset
from Yoffe import regularizedYoffe
from scipy import ndimage

def writeNetcdf(sname, lDimVar, lName, lData, paraview_readable=False):
    """create a netcdf file either readable by ASAGI
    or by paraview (paraview_readable=True)

    Parameters
    ----------
    sname: str prefix name of output file
    lDimVar: list of 1d numpy array containing the dimension variables
    lName: list if str containing the name of the nd variables
    lData: list if n-d numpy array containing the data
    paraview_readable: bool
    """
    fname = f"{sname}.nc"
    print("writing " + fname)

    with Dataset(fname, "w", format="NETCDF4") as rootgrp:
        # Create dimension and 1d variables
        sdimVarNames = "uvwxyz"
        dims = []
        for i, xi in enumerate(lDimVar):
            nxi = xi.shape[0]
            dimName = sdimVarNames[i]
            dims.append(dimName)
            rootgrp.createDimension(dimName, nxi)
            vx = rootgrp.createVariable(dimName, "f4", (dimName,))
            vx[:] = xi
        dims.reverse()
        dims = tuple(dims)

        if paraview_readable:
            for i in range(len(lName)):
                vTd = rootgrp.createVariable(lName[i], "f4", dims)
                vTd[:] = lData[i]
        else:
            ldata4 = [(name, "f4") for name in lName]
            ldata8 = [(name, "f8") for name in lName]
            mattype4 = np.dtype(ldata4)
            mattype8 = np.dtype(ldata8)
            mat_t = rootgrp.createCompoundType(mattype4, "material")

            # this transform the nD array into an array of tuples
            arr = np.stack([lData[i] for i in range(len(lName))], axis=len(dims))
            newarr = arr.view(dtype=mattype8)
            newarr = newarr.reshape(newarr.shape[:-1])
            mat = rootgrp.createVariable("data", mat_t, dims)
            mat[:] = newarr

def cosine_taper(npts, p=0.1, freqs=None, flimit=None, halfcosine=True, sactaper=False):
    """
    Cosine Taper. (copied from obspy:
    https://docs.obspy.org/master/_modules/obspy/signal/invsim.html#cosine_taper

    :type npts: int
    :param npts: Number of points of cosine taper.
    :type p: float
    :param p: Decimal percentage of cosine taper (ranging from 0 to 1). Default
        is 0.1 (10%) which tapers 5% from the beginning and 5% form the end.
    :rtype: float NumPy :class:`~numpy.ndarray`
    :return: Cosine taper array/vector of length npts.
    :type freqs: NumPy :class:`~numpy.ndarray`
    :param freqs: Frequencies as, for example, returned by fftfreq
    :type flimit: list or tuple of floats
    :param flimit: The list or tuple defines the four corner frequencies
        (f1, f2, f3, f4) of the cosine taper which is one between f2 and f3 and
        tapers to zero for f1 < f < f2 and f3 < f < f4.
    :type halfcosine: bool
    :param halfcosine: If True the taper is a half cosine function. If False it
        is a quarter cosine function.
    :type sactaper: bool
    :param sactaper: If set to True the cosine taper already tapers at the
        corner frequency (SAC behavior). By default, the taper has a value
        of 1.0 at the corner frequencies.

    .. rubric:: Example

    >>> tap = cosine_taper(100, 1.0)
    >>> tap2 = 0.5 * (1 + np.cos(np.linspace(np.pi, 2 * np.pi, 50)))
    >>> np.allclose(tap[0:50], tap2)
    True
    >>> npts = 100
    >>> p = 0.1
    >>> tap3 = cosine_taper(npts, p)
    >>> (tap3[int(npts*p/2):int(npts*(1-p/2))]==np.ones(int(npts*(1-p)))).all()
    True
    """
    if p < 0 or p > 1:
        msg = "Decimal taper percentage must be between 0 and 1."
        raise ValueError(msg)
    if p == 0.0 or p == 1.0:
        frac = int(npts * p / 2.0)
    else:
        frac = int(npts * p / 2.0 + 0.5)

    if freqs is not None and flimit is not None:
        fl1, fl2, fl3, fl4 = flimit
        idx1 = np.argmin(abs(freqs - fl1))
        idx2 = np.argmin(abs(freqs - fl2))
        idx3 = np.argmin(abs(freqs - fl3))
        idx4 = np.argmin(abs(freqs - fl4))
    else:
        idx1 = 0
        idx2 = frac - 1
        idx3 = npts - frac
        idx4 = npts - 1
    if sactaper:
        # in SAC the second and third
        # index are already tapered
        idx2 += 1
        idx3 -= 1

    # Very small data lengths or small decimal taper percentages can result in
    # idx1 == idx2 and idx3 == idx4. This breaks the following calculations.
    if idx1 == idx2:
        idx2 += 1
    if idx3 == idx4:
        idx3 -= 1

    # the taper at idx1 and idx4 equals zero and
    # at idx2 and idx3 equals one
    cos_win = np.zeros(npts)
    if halfcosine:
        # cos_win[idx1:idx2+1] =  0.5 * (1.0 + np.cos((np.pi * \
        #    (idx2 - np.arange(idx1, idx2+1)) / (idx2 - idx1))))
        cos_win[idx1 : idx2 + 1] = 0.5 * (1.0 - np.cos((np.pi * (np.arange(idx1, idx2 + 1) - float(idx1)) / (idx2 - idx1))))
        cos_win[idx2 + 1 : idx3] = 1.0
        cos_win[idx3 : idx4 + 1] = 0.5 * (1.0 + np.cos((np.pi * (float(idx3) - np.arange(idx3, idx4 + 1)) / (idx4 - idx3))))
    else:
        cos_win[idx1 : idx2 + 1] = np.cos(-(np.pi / 2.0 * (float(idx2) - np.arange(idx1, idx2 + 1)) / (idx2 - idx1)))
        cos_win[idx2 + 1 : idx3] = 1.0
        cos_win[idx3 : idx4 + 1] = np.cos((np.pi / 2.0 * (float(idx3) - np.arange(idx3, idx4 + 1)) / (idx4 - idx3)))

    # if indices are identical division by zero
    # causes NaN values in cos_win
    if idx1 == idx2:
        cos_win[idx1] = 0.0
    if idx3 == idx4:
        cos_win[idx3] = 0.0
    return cos_win

def interpolate_nan_from_neighbors(array):
    """rise_time and tacc may not be defined where there is no slip (no SR function).
    in this case, we interpolate from neighbors
    source: https://stackoverflow.com/questions/37662180/interpolate-missing-values-2d-python
    """
    x = np.arange(0, array.shape[1])
    y = np.arange(0, array.shape[0])
    # mask invalid values
    array = np.ma.masked_invalid(array)
    xx, yy = np.meshgrid(x, y)
    # get only the valid values
    x1 = xx[~array.mask]
    y1 = yy[~array.mask]
    newarr = array[~array.mask]
    return interpolate.griddata((x1, y1), newarr.ravel(), (xx, yy), method="linear", fill_value=np.average(array))

def compute_block_mean(ar, fact):
    """
    dowsample array ar by factor fact
    https://stackoverflow.com/questions/18666014/downsample-array-in-python
    """
    assert isinstance(fact, int), type(fact)
    sx, sy = ar.shape
    X, Y = np.ogrid[0:sx, 0:sy]
    regions = sy // fact * (X // fact) + Y // fact
    res = ndimage.mean(ar, labels=regions, index=np.arange(regions.max() + 1))
    res.shape = (sx // fact, sy // fact)
    return res

def upsample_quantities(allarr, spatial_order, spatial_zoom, padding="constant", extra_padding_layer=False, minimize_block_average_variations=False):
    """1. pad
    2. upsample, adding spatial_zoom per node
    """
    nd = allarr.shape[0]
    ny, nx = [val * spatial_zoom for val in allarr[0].shape]
    if extra_padding_layer:
        # required for vertex aligned netcdf format
        nx = nx + 2
        ny = ny + 2
    allarr0 = np.zeros((nd, ny, nx))
    for k in range(nd):
        if padding == "extrapolate":
            my_array0 = np.pad(allarr[k, :, :], ((1, 1), (1, 1)), "reflect", reflect_type="odd")
        else:
            my_array0 = np.pad(allarr[k, :, :], ((1, 1), (1, 1)), padding)
        if extra_padding_layer:
            ncrop = spatial_zoom - 1
        else:
            ncrop = spatial_zoom
        my_array = scipy.ndimage.zoom(my_array0, spatial_zoom, order=spatial_order, mode="grid-constant", grid_mode=True)
        if minimize_block_average_variations:
            # inspired by Tinti et al. (2005) (Appendix A)
            # This is for the specific case of fault slip.
            # We want to preserve the seismic moment of each subfault after interpolation
            # the rock rigidity is not know by this script (would require some python binding of easi).
            # the subfault area is typically constant over the kinematic model
            # So we just want to perserve subfault average.
            print("trying to perserve subfault average...")
            my_array = np.maximum(0, my_array)
            best_misfit = float("inf")
            # The algorithm does not seem to converge, but produces better model
            # (given the misfit) that inital after 2-3 iterations
            niter = 30
            for i in range(niter):
                block_average = compute_block_mean(my_array, spatial_zoom)
                correction = my_array0 / block_average
                # having a misfit as misfit = np.linalg.norm(correction) does not makes sense as for almost 0 slip, correction can be large
                misfit = np.linalg.norm(my_array0 - block_average) / len(my_array0)
                if best_misfit > misfit:
                    if i == 0:
                        print(f"misfit at iter {i}: {misfit}")
                    else:
                        print(f"misfit improved at iter {i}: {misfit}")
                    best_misfit = misfit
                    best = np.copy(my_array)
                my_array = scipy.ndimage.zoom(correction * my_array0, spatial_zoom, order=spatial_order, mode="grid-constant", grid_mode=True)
                my_array = np.maximum(0, my_array)
            my_array = best
        if ncrop > 0:
            allarr0[k, :, :] = my_array[ncrop:-ncrop, ncrop:-ncrop]

    return allarr0

class FaultPlane:
    def __init__(self):
        self.nx = 0
        self.ny = 0
        self.ndt = 0
        self.PSarea_cm2 = 0
        self.dt = 0
        # array member initialized to dummy value
        self.lon = 0
        self.lat = 0
        self.x = 0
        self.y = 0
        self.depth = 0
        self.t0 = 0
        self.slip1 = 0
        self.strike = 0
        self.dip = 0
        self.rake = 0
        self.aSR = 0
        self.myt = 0

    def init_spatial_arrays(self, nx, ny):
        self.nx = nx
        self.ny = ny
        self.lon = np.zeros((ny, nx))
        self.lat = np.zeros((ny, nx))
        self.x = np.zeros((ny, nx))
        self.y = np.zeros((ny, nx))
        self.depth = np.zeros((ny, nx))
        self.t0 = np.zeros((ny, nx))
        self.slip1 = np.zeros((ny, nx))
        self.strike = np.zeros((ny, nx))
        self.dip = np.zeros((ny, nx))
        self.rake = np.zeros((ny, nx))

    def init_aSR(self):
        self.aSR = np.zeros((self.ny, self.nx, self.ndt))

    def extend_aSR(self, ndt_old, ndt_new):
        "extend aSR array to more time samplings"
        tmpSR = np.copy(self.aSR)
        self.ndt = ndt_new
        self.aSR = np.zeros((self.ny, self.nx, self.ndt))
        self.aSR[:, :, 0:ndt_old] = tmpSR[:, :, :]

    def compute_xy_from_latlon(self, proj):
        if proj:
            from pyproj import Transformer

            transformer = Transformer.from_crs("epsg:4326", proj[0], always_xy=True)
            self.x, self.y = transformer.transform(self.lon, self.lat)
        else:
            print("no proj string specified!")
            self.x, self.y = self.lon, self.lat

    def compute_latlon_from_xy(self, proj):
        if proj:
            from pyproj import Transformer

            transformer = Transformer.from_crs(proj[0], "epsg:4326", always_xy=True)
            self.lon, self.lat = transformer.transform(self.x, self.y)
        else:
            self.lon, self.lat = self.x, self.y

    def compute_time_array(self):
        self.myt = np.linspace(0, (self.ndt - 1) * self.dt, self.ndt)

    def write_srf(self, fname):
        "write kinematic model to a srf file (standard rutpure format)"
        with open(fname, "w") as fout:
            fout.write("1.0\n")
            fout.write("POINTS %d\n" % (self.nx * self.ny))
            for j in range(self.ny):
                for i in range(self.nx):
                    fout.write("%g %g %g %g %g %e %g %g\n" % (self.lon[j, i], self.lat[j, i], self.depth[j, i], self.strike[j, i], self.dip[j, i], self.PSarea_cm2, self.t0[j, i], self.dt))
                    fout.write("%g %g %d %f %d %f %d\n" % (self.rake[j, i], self.slip1[j, i], self.ndt, 0.0, 0, 0.0, 0))
                    np.savetxt(fout, self.aSR[j, i, :], fmt="%g", newline=" ")
                    fout.write("\n")
        print("done writing", fname)

    def init_from_srf(self, fname):
        "init object by reading a srf file (standard rutpure format)"
        with open(fname) as fid:
            # version
            line = fid.readline()
            version = float(line)
            if not (abs(version - 1.0) < 1e-03 or abs(version - 2.0) < 1e-03):
                print("srf version: %s not supported" % (line))
                raise
            # skip comments
            while True:
                line = fid.readline()
                if not line.startswith("#"):
                    break
            line_el = line.split()
            if line_el[0] != "PLANE":
                print("no plane specified")
                raise
            if line_el[1] != "1":
                print("only one plane supported")
                raise NotImplementedError
            line_el = fid.readline().split()
            nx, ny = [int(val) for val in line_el[2:4]]
            line_el = fid.readline().split()
            # check that the plane data are consistent with the number of points
            assert int(line_el[1]) == nx * ny
            self.init_spatial_arrays(nx, ny)
            for j in range(ny):
                for i in range(nx):
                    # first header line
                    line = fid.readline()
                    # rho_vs are only present for srf version 2
                    self.lon[j, i], self.lat[j, i], self.depth[j, i], self.strike[j, i], self.dip[j, i], self.PSarea_cm2, self.t0[j, i], dt, *rho_vs = [float(v) for v in line.split()]
                    # second header line
                    line = fid.readline()
                    self.rake[j, i], self.slip1[j, i], ndt1, slip2, ndt2, slip3, ndt3 = [float(v) for v in line.split()]
                    if max(slip2, slip3) > 0.0:
                        print("this script assumes slip2 and slip3 are zero", slip2, slip3)
                        raise NotImplementedError
                    ndt1 = int(ndt1)
                    if max(i, j) == 0:
                        self.ndt = ndt1
                        self.dt = dt
                        self.init_aSR()
                    lSTF = []
                    if ndt1 == 0:
                        continue
                    if ndt1 > self.ndt:
                        print(f"a larger ndt ({ndt1}> {self.ndt}) was found for point source (i,j) = ({i}, {j}) extending aSR array...")
                        self.extend_aSR(self.ndt, ndt1)
                    if abs(dt - self.dt) > 1e-6:
                        print("this script assumes that dt is the same for all sources", dt, self.dt)
                        raise NotImplementedError
                    while True:
                        line = fid.readline()
                        lSTF.extend(line.split())
                        if len(lSTF) == ndt1:
                            self.aSR[j, i, 0:ndt1] = np.array([float(v) for v in lSTF])
                            break

    def assess_STF_parameters(self):
        "compute rise_time (slip duration) and t_acc (peak SR) from SR time histories"
        self.rise_time = np.zeros((self.ny, self.nx))
        self.tacc = np.zeros((self.ny, self.nx))
        for j in range(self.ny):
            for i in range(self.nx):
                if not self.slip1[j, i]:
                    self.rise_time[j, i] = np.nan
                    self.tacc[j, i] = np.nan
                else:
                    first_non_zero = np.amin(np.where(self.aSR[j, i, :])[0])
                    last_non_zero = np.amax(np.where(self.aSR[j, i, :])[0])
                    id_max = np.where(self.aSR[j, i, :] == np.amax(self.aSR[j, i, :]))[0]
                    self.rise_time[j, i] = (last_non_zero - first_non_zero + 1) * self.dt
                    self.tacc[j, i] = (id_max - first_non_zero + 1) * self.dt
                    self.t0[j, i] += first_non_zero * self.dt
        self.rise_time = interpolate_nan_from_neighbors(self.rise_time)
        self.tacc = interpolate_nan_from_neighbors(self.tacc)

        print("slip rise_time (min, 50%, max)", np.amin(self.rise_time), np.median(self.rise_time), np.amax(self.rise_time))
        print("tacc (min, 50%, max)", np.amin(self.tacc), np.median(self.tacc), np.amax(self.tacc))

    def upsample_fault(self, spatial_order, spatial_zoom, temporal_zoom, proj, use_Yoffe=False, time_smoothing_kernel_as_dt_fraction=0.5):
        "increase spatial and temporal resolution of kinematic model by interpolation"
        # time vector
        ndt2 = (self.ndt - 1) * temporal_zoom + 1
        ny2, nx2 = self.ny * spatial_zoom, self.nx * spatial_zoom
        # resampled source
        pf = FaultPlane()
        pf.init_spatial_arrays(nx2, ny2)
        pf.ndt = ndt2
        pf.init_aSR()

        pf.dt = self.dt / temporal_zoom
        pf.compute_time_array()

        # upsample spatially geometry (bilinear interpolation)
        allarr = np.array([self.x, self.y, self.depth])
        pf.x, pf.y, pf.depth = upsample_quantities(allarr, spatial_order=1, spatial_zoom=spatial_zoom, padding="extrapolate")

        # upsample other quantities
        allarr = np.array([self.t0, self.strike, self.dip, self.rake])
        pf.t0, pf.strike, pf.dip, pf.rake = upsample_quantities(allarr, spatial_order, spatial_zoom, padding="edge")
        # the interpolation may generate some acausality that we here prevent
        pf.t0 = np.maximum(pf.t0, np.amin(self.t0))

        allarr = np.array([self.slip1])
        (pf.slip1,) = upsample_quantities(allarr, spatial_order, spatial_zoom, padding="constant", minimize_block_average_variations=True)
        pf.compute_latlon_from_xy(proj)
        pf.PSarea_cm2 = self.PSarea_cm2 / spatial_zoom ** 2
        ratio_potency = np.sum(pf.slip1) * pf.PSarea_cm2 / (np.sum(self.slip1) * self.PSarea_cm2)
        print(f"seismic potency ratio (upscaled over initial): {ratio_potency}")

        if use_Yoffe:
            self.assess_STF_parameters()
            allarr = np.array([self.rise_time, self.tacc])
            pf.rise_time, pf.tacc = upsample_quantities(allarr, spatial_order, spatial_zoom, padding="edge")
            pf.rise_time = np.maximum(pf.rise_time, np.amin(self.rise_time))
            pf.tacc = np.maximum(pf.tacc, np.amin(self.tacc))
            print("using ts = tacc / 1.27 to compute the regularized Yoffe")
            ts = pf.tacc / 1.27
            tr = pf.rise_time - 2.0 * ts
            for j in range(pf.ny):
                for i in range(pf.nx):
                    for k, tk in enumerate(pf.myt):
                        pf.aSR[j, i, k] = pf.slip1[j, i] * regularizedYoffe(tk, ts[j, i], tr[j, i])
        else:
            aSRa = np.zeros((pf.ny, pf.nx, self.ndt))
            for k in range(self.ndt):
                aSRa[:, :, k] = upsample_quantities(np.array([self.aSR[:, :, k]]), spatial_order, spatial_zoom, padding="constant")

            # interpolate temporally the AST
            for j in range(pf.ny):
                for i in range(pf.nx):
                    # 1. upsample with linear interpolation
                    # 2. apply a gauss kernel to smooth out sharp edges
                    # 3. tapper the signal smoothly to 0 at both time ends
                    # 4. rescale SR to ensure integral (SR) = slip
                    f = interpolate.interp1d(self.myt, aSRa[j, i, :], kind="linear")
                    pf.aSR[j, i, :] = f(pf.myt)
                    tapper = cosine_taper(pf.ndt, self.dt / (pf.ndt * pf.dt))
                    pf.aSR[j, i, :] = tapper * ndimage.gaussian_filter1d(pf.aSR[j, i, :], time_smoothing_kernel_as_dt_fraction * self.dt / pf.dt, mode="constant")
                    # With a cubic interpolation, the interpolated slip1 may be negative which does not make sense.
                    if pf.slip1[j, i] < 0:
                        pf.aSR[j, i, :] = 0
                        continue
                    # should be the SR
                    integral_STF = np.trapz(np.abs(pf.aSR[j, i, :]), dx=pf.dt)
                    if abs(integral_STF) > 0:
                        pf.aSR[j, i, :] = pf.slip1[j, i] * pf.aSR[j, i, :] / integral_STF
        return pf

    def compute_corrected_slip_for_differing_area(self, proj):
        """
        self.PSarea_cm2 may slightly differ from the patch area from the fault geometry
        (e.g. due to the projection)
        Therefore, we need to update slip to keep seismic potency (area*slip) unchanged
        """
        cm2m = 0.01
        km2m = 1e3
        PSarea_m2 = self.PSarea_cm2 * cm2m * cm2m
        self.compute_xy_from_latlon(proj)
        nx, ny = self.nx, self.ny
        # Compute actual dx and dy from coordinates
        dy = np.zeros((ny, nx))
        dx = np.zeros((ny, nx))
        # central difference for the inside
        coords = np.array((self.x, self.y, -km2m * self.depth))
        for i in range(0, nx):
            p0 = coords[:, 0 : ny - 2, i] - coords[:, 2:ny, i]
            dy[1 : ny - 1, i] = 0.5 * np.linalg.norm(p0, axis=0)
        # special case of 0 and ny-1
        p0 = coords[:, 1, :] - coords[:, 0, :]
        dy[0, :] = np.linalg.norm(p0, axis=0)
        p0 = coords[:, ny - 1, :] - coords[:, ny - 2, :]
        dy[ny - 1, :] = np.linalg.norm(p0, axis=0)
        # dx for coordinates
        for j in range(0, ny):
            p0 = coords[:, j, 0 : nx - 2] - coords[:, j, 2:nx]
            dx[j, 1 : nx - 1] = 0.5 * np.linalg.norm(p0, axis=0)
        p0 = coords[:, :, 1] - coords[:, :, 0]
        dx[:, 0] = np.linalg.norm(p0, axis=0)
        p0 = coords[:, :, nx - 1] - coords[:, :, nx - 2]
        dx[:, nx - 1] = np.linalg.norm(p0, axis=0)
        factor_area = dx[:, :] * dy[:, :] / PSarea_m2
        slip1 = self.slip1 * factor_area
        print(
            f"done correcting slip for area. \
The correcting factor ranges between {np.amin(factor_area)} and {np.amax(factor_area)}"
        )
        return slip1

    def generate_netcdf_fl33(self, prefix, spatial_order, spatial_zoom, proj, write_paraview):
        "generate netcdf files to be used with SeisSol friction law 33"

        cm2m = 0.01
        km2m = 1e3
        # a kinematic model defines the fault quantities at the subfault center
        # a netcdf file defines the quantities at the nodes
        # therefore the extra_padding_layer=True, and the added di below
        cslip = self.compute_corrected_slip_for_differing_area(proj)
        (slip,) = upsample_quantities(np.array([cslip]), spatial_order, spatial_zoom, padding="constant", extra_padding_layer=True, minimize_block_average_variations=True)
        allarr = np.array([self.t0, self.rake, self.rise_time, self.tacc])
        rupttime, rake, rise_time, tacc = upsample_quantities(allarr, spatial_order, spatial_zoom, padding="edge", extra_padding_layer=True)
        # upsampled duration, rise_time and acc_time may not be smaller than initial values
        # at least rise_time could lead to a non-causal kinematic model
        rupttime = np.maximum(rupttime, np.amin(self.t0))
        rise_time = np.maximum(rise_time, np.amin(self.rise_time))
        tacc = np.maximum(tacc, np.amin(self.tacc))

        rake_rad = np.radians(rake)
        strike_slip = slip * np.cos(rake_rad) * cm2m
        dip_slip = slip * np.sin(rake_rad) * cm2m

        ny, nx = slip.shape
        dx = np.sqrt(self.PSarea_cm2 * cm2m * cm2m)
        ldataName = ["strike_slip", "dip_slip", "rupture_onset", "effective_rise_time", "acc_time"]
        lgridded_myData = [strike_slip, dip_slip, rupttime, rise_time, tacc]
        # we could do directly
        # di = 1.0 / (2 * spatial_zoom)
        # xb = np.linspace(-di, self.nx + di, nx) * dx
        # yb = np.linspace(-di, self.ny + di, ny) * dx
        # But we compute xb and yb based on the upsampled coordinates, to account for possible warping due to the projection
        allarr = np.array([self.x, self.y, -km2m * self.depth])
        coords = upsample_quantities(allarr, spatial_order=1, spatial_zoom=spatial_zoom, padding="extrapolate", extra_padding_layer=True)

        p0 = coords[:, (ny - 1) // 2, :] - coords[:, (ny - 1) // 2 - 1, :]
        dx1 = np.linalg.norm(p0, axis=0)
        xb = np.cumsum(dx1) - 1.5 * dx1[0]

        p0 = coords[:, :, (nx - 1) // 2] - coords[:, :, (nx - 1) // 2 - 1]
        dy1 = np.linalg.norm(p0, axis=0)
        yb = np.cumsum(dy1) - 1.5 * dy1[0]

        prefix2 = f"{prefix}_{spatial_zoom}_o{spatial_order}"
        if write_paraview:
            # see comment above
            for i, sdata in enumerate(ldataName):
                writeNetcdf(prefix2 + sdata, [xb, yb], [sdata], [lgridded_myData[i]], paraview_readable=True)
        writeNetcdf(prefix2, [xb, yb], ldataName, lgridded_myData)

    def generate_fault_ts_yaml_fl33(self, prefix, spatial_order, spatial_zoom, proj):
        """Generate yaml file initializing FL33 arrays and ts file describing the planar fault geometry."""
        # Generate yaml file loading ASAGI file
        cm2m = 0.01
        km2m = 1e3
        self.compute_xy_from_latlon(proj)
        nx, ny = self.nx, self.ny
        p0 = np.array([self.x[0, 0], self.y[0, 0], -km2m * self.depth[0, 0]])
        p1 = np.array([self.x[ny - 1, 0], self.y[ny - 1, 0], -km2m * self.depth[ny - 1, 0]])
        p2 = np.array([self.x[0, nx - 1], self.y[0, nx - 1], -km2m * self.depth[0, nx - 1]])
        p3 = np.array([self.x[ny - 1, nx - 1], self.y[ny - 1, nx - 1], -km2m * self.depth[ny - 1, nx - 1]])

        hw = p1 - p0
        dx1 = np.linalg.norm(hw) / (ny - 1)
        hw = hw / np.linalg.norm(hw)
        hh = p2 - p0
        dx2 = np.linalg.norm(hh) / (nx - 1)
        hh = hh / np.linalg.norm(hh)
        dx = np.sqrt(self.PSarea_cm2 * cm2m * cm2m)
        # a kinematic model defines the fault quantities at the subfault center
        # a netcdf file defines the quantities at the nodes
        # therefore the dx/2
        # the term dxi/np.sqrt(dx1*dx2) allows accounting for non-square patches
        non_square_factor = dx / np.sqrt(dx1 * dx2)
        t1 = -np.dot(p0, hh) + 0.5 * dx1 * non_square_factor
        t2 = -np.dot(p0, hw) + 0.5 * dx2 * non_square_factor

        template_yaml = f"""!Switch
[strike_slip, dip_slip, rupture_onset, tau_S, tau_R, rupture_rise_time]: !EvalModel
    parameters: [strike_slip, dip_slip, rupture_onset, effective_rise_time, acc_time]
    model: !Switch
        [strike_slip, dip_slip, rupture_onset, effective_rise_time, acc_time]: !AffineMap
              matrix:
                ua: [{hh[0]}, {hh[1]}, {hh[2]}]
                ub: [{hw[0]}, {hw[1]}, {hw[2]}]
              translation:
                ua: {t1}
                ub: {t2}
              components: !Any
                - !ASAGI
                    file: {prefix}_{spatial_zoom}_o{spatial_order}.nc
                    parameters: [strike_slip, dip_slip, rupture_onset, effective_rise_time, acc_time]
                    var: data
                    interpolation: linear
                - !ConstantMap
                  map:
                    strike_slip: 0.0
                    dip_slip:    0.0
                    rupture_onset:    0.0
                    acc_time:  1e100
                    effective_rise_time:  2e100
    components: !FunctionMap
       map:
          #Note the minus on strike_slip to acknowledge the different convention of SeisSol (T_s>0 means right-lateral)
          strike_slip: return -strike_slip;
          dip_slip: return dip_slip;
          rupture_onset: return rupture_onset;
          tau_S: return acc_time/1.27;
          tau_R: return effective_rise_time - 2.*acc_time/1.27;
          rupture_rise_time: return effective_rise_time;
        """
        fname = f"{prefix}_fault.yaml"
        with open(fname, "w") as fid:
            fid.write(template_yaml)
        print(f"done writing {fname}")

        # Generate ts file containing mesh geometry
        vertex = np.zeros((4, 3))
        vertex[0, :] = p0 + 0.5 * (-hh * dx1 - hw * dx2) * non_square_factor
        vertex[1, :] = p2 + 0.5 * (hh * dx1 - hw * dx2) * non_square_factor
        vertex[2, :] = p3 + 0.5 * (hh * dx1 + hw * dx2) * non_square_factor
        vertex[3, :] = p1 + 0.5 * (-hh * dx1 + hw * dx2) * non_square_factor

        connect = np.zeros((2, 3), dtype=int)
        connect[0, :] = [1, 2, 3]
        connect[1, :] = [1, 3, 4]
        fname = f"{prefix}_fault.ts"
        with open(fname, "w") as fout:
            fout.write("GOCAD TSURF 1\nHEADER {\nname:%s\nborder: true\nmesh: false\n*border*bstone: true\n}\nTFACE\n" % (fname))
            for ivx in range(1, 5):
                fout.write("VRTX %s %s %s %s\n" % (ivx, vertex[ivx - 1, 0], vertex[ivx - 1, 1], vertex[ivx - 1, 2]))

            for i in range(2):
                fout.write("TRGL %d %d %d\n" % (connect[i, 0], connect[i, 1], connect[i, 2]))
            fout.write("END\n")
        print(f"done writing {fname}")


import logging
import os
from typing import Dict, List, Optional

import numpy as np
import scipy.fft as fp
from scipy.ndimage import generic_gradient_magnitude, sobel
from skimage.transform import resize

import nanotune as nt

logger = logging.getLogger(__name__)
N_2D = nt.config["core"]["standard_shapes"]["2"]
NT_LABELS = list(dict(nt.config["core"]["labels"]).keys())

def export_label(
    ds_label: List["str"],
    quality: int,
    category: str,
) -> int:
    """Merges binary labels of single and double dot qualities to a single
    label. Only if `dotregime` is specified, for all other the initial
    quality labels are returned.
    It translates a dot regime label to: 0 - poor singledot, 1 - good singledot,
    2 - poor doubledot, 3 - good doubledot.

    Returns:
        int: new label.
    """
    good = bool(quality)
    if category == "dotregime":
        singledot = True if "singledot" in ds_label else False
        doubledot = True if "doubledot" in ds_label else False

        if not good and singledot:
            new_label = 0
        if good and singledot:
            new_label = 1
        if not good and doubledot:
            new_label = 2
        if good and doubledot:
            new_label = 3

    elif category in ["outerbarriers", "pinchoff", "singledot", "doubledot"]:
        if len(ds_label) == 1 and ds_label[0] == category:
            new_label = good
        else:
            print(category)
            print(ds_label)
            logger.warning("Wrong label-category combination in export_label.")
            raise ValueError
    else:
        logger.error(
            "Do not know how to export/condense labels. Please "
            + "update export_label in export_data.py."
        )
        raise ValueError

    return new_label

def prep_data(
    dataset: nt.Dataset,
    category: str,
    flip_data: bool = False,
    readout_method_to_use: str = 'transport',
) -> List[List[List[float]]]:
    """Prepares data for classification.
    It combines normalized data, its gradient, Fourier frequencies and
    extracted features into a multidimensional list.
    All sublists are reshaped to a standard shape, defined in config.json
    under the `standard_shapes` key.

    Args:
        dataset: instance of nanotune dataset whose data should be prepared.
        category: as which category/type of data, e.g. `pinchoff`, `singledot`
            etc, it should be treated.
        flip_data: whether data should be flipped. Used to simulate pinchoff
            curves measured in rf sensing.

    Returns:
        list: multidimensional list. First sublist is normalized data,
            second Fourier frequencies, third gradient and fourth features.
    """
    assert category in nt.config["core"]["features"].keys()
    if len(dataset.power_spectrum) == 0:
        dataset.compute_power_spectrum()

    condensed_data_all = []

    # for readout_method in dataset.readout_methods.keys():
    signal = dataset.data[readout_method_to_use].values
    if flip_data:
        signal = np.flip(signal)
    dimension = dataset.dimensions[readout_method_to_use]

    shape = tuple(nt.config["core"]["standard_shapes"][str(dimension)])
    condensed_data = np.empty(
        (len(nt.config["core"]["data_types"]), 1, np.prod(shape))
    )

    relevant_features = nt.config["core"]["features"][category]
    features = []

    if dataset.features:
        if all(isinstance(i, dict) for i in dataset.features.values()):
            for feat in relevant_features:
                features.append(dataset.features[readout_method_to_use][feat])
        else:
            for feat in relevant_features:
                features.append(dataset.features[feat])

    # double check if current range is correct:
    if np.max(signal) > 1:
        min_curr = np.min(signal)
        max_curr = np.max(signal)
        signal = (signal - min_curr) / (max_curr - min_curr)
        # assume we are talking dots and high current was not actually
        # device_max_signal
        dataset.data[readout_method_to_use].values = signal * 0.3
        dataset.compute_power_spectrum()

    data_resized = resize(
        signal, shape, anti_aliasing=True, mode="edge"
    ).flatten()

    grad = generic_gradient_magnitude(signal, sobel)
    gradient_resized = resize(
        grad, shape, anti_aliasing=True, mode="constant"
    ).flatten()
    power = dataset.power_spectrum[readout_method_to_use].values
    frequencies_resized = resize(
        power, shape, anti_aliasing=True, mode="constant"
    ).flatten()

    pad_width = len(data_resized.flatten()) - len(features)
    features = np.pad(
        features,
        (0, pad_width),
        "constant",
        constant_values=nt.config["core"]["fill_value"],
    )

    index = nt.config["core"]["data_types"]["signal"]
    condensed_data[index, 0, :] = data_resized.tolist()

    index = nt.config["core"]["data_types"]["frequencies"]
    condensed_data[index, 0, :] = frequencies_resized.tolist()

    index = nt.config["core"]["data_types"]["gradient"]
    condensed_data[index, 0, :] = gradient_resized.tolist()

    index = nt.config["core"]["data_types"]["features"]
    condensed_data[index, 0, :] = features

    condensed_data_all.append(condensed_data.tolist())

    return condensed_data_all

def export_data(
    category: str,
    db_names: List[str],
    skip_ids: Optional[Dict[str, List[int]]] = None,
    add_flipped_data: bool = False,
    quality: Optional[int] = None,
    filename: Optional[str] = None,
    db_folder: Optional[str] = None,
    readout_method_to_use: str = 'transport',
) -> None:
    """Exports condensed data to a numpy file in a format used by
    nanotune's Classifier.

    The saved array contains the normalized signal, gradient, Fourier
    frequencies and extracted features of each dataset. The machine learning
    labels are attached to each of them to the very end.
    A dataset of 15 1D measurements results in an array of (4, 15, 101),
    where each trace has been reshaped to 100 points plus the label. The
    first array[0, :, :] are all signals, array[1, :, :] the frequencies,
    array is defined in config.json under `data_types`.

    Args:
        category: nt.config['core']['features'].keys()
        db_names: names of databases whose data should be exported.
        skip_ids: dict mapping database names to list of run IDs which should
            be skipped.
        add_flipped_data: whether or flipped data should be added as well.
            a flipped pinchoff curve measured in transport for example
            reproduces a pinchoff measured in rf.
        quality: which qualities, e.g. only good or poor data, should be
            exported. If none give, both will be taken.
        filename: name of resulting numpy file. Default is the category name.
        db_folder: folder where databases are located.
        readout_method_to_use: which readout method to use if more than one
            is available. Default is 'transport'.
    """
    assert isinstance(db_names, list)
    if category not in list(nt.config['core']['features'].keys()):
        raise ValueError(
            f"Unknown category. Please use on of the following: \
            {list(nt.config['core']['features'].keys())}."
        )

    if db_folder is None:
        db_folder = nt.config["db_folder"]

    if category in ["pinchoff", "coulomboscillation", "zerobiaspeak1D"]:
        dim = 1
    else:
        dim = 2

    if category == 'dotregime':
        stages = ["singledot", "doubledot"]
    else:
        stages = [category]

    shape = tuple(nt.config["core"]["standard_shapes"][str(dim)])
    condensed_data_all = np.empty(
        (len(nt.config["core"]["data_types"]), 0, np.prod(shape))
    )

    relevant_ids: Dict[str, List[int]] = {}
    for db_name in db_names:
        relevant_ids[db_name] = []
        nt.set_database(db_name, db_folder)
        for stage in stages:
            try:
                relevant_ids[db_name] += nt.get_dataIDs(
                    db_name, stage, quality=quality, db_folder=db_folder,
                    get_run_ids=False,
                )
                if len(relevant_ids[db_name]) == 0:
                    logger.warning(f'No labelled data found in {db_name}')
            except Exception as e:
                msg = f"Unable to load relevant ids in {db_name}." + str(e)
                logger.error(msg)
                break

    labels_exp = []
    for db_name, dataids in relevant_ids.items():
        nt.set_database(db_name, db_folder)
        skip_us = []
        if skip_ids is not None:
            try:
                skip_us = skip_ids[db_name]
            except KeyError:
                logger.warning("No data IDs to skip in {}.".format(db_name))

        for d_id in dataids:
            if d_id not in skip_us:
                try:
                    df = nt.Dataset(d_id, db_name, db_folder=db_folder)
                    condensed_data = prep_data(
                        df,
                        category,
                        readout_method_to_use=readout_method_to_use)
                    new_label = export_label(df.ml_label, df.quality, category)
                    condensed_data_all = np.append(
                        condensed_data_all, condensed_data[0], axis=1
                    )
                    labels_exp.append(new_label)

                    if add_flipped_data:
                        condensed_data = prep_data(
                            df, category, flip_data=True,
                            readout_method_to_use=readout_method_to_use
                        )
                        new_label = export_label(
                            df.ml_label, df.quality, category
                        )
                        condensed_data_all = np.append(
                            condensed_data_all, condensed_data[0], axis=1
                        )
                        labels_exp.append(new_label)
                except (IndexError, ValueError, TypeError) as i_err:
                    print(db_name)
                    print(d_id)
                    print(i_err)

    n = list(condensed_data_all.shape)
    n[-1] += 1

    data_w_labels = np.zeros(n)
    data_w_labels[:, :, -1] = labels_exp
    data_w_labels[:, :, :-1] = condensed_data_all

    if filename is None:
        filename = "_".join(stages)
    path = os.path.join(db_folder, filename)
    np.save(path, data_w_labels)

def correct_normalizations(
    filename: str,
    db_folder: Optional[str] = None,
) -> None:
    """"""
    if db_folder is None:
        db_folder = nt.config["db_folder"]

    path = os.path.join(db_folder, filename)

    all_data = np.load(path)

    data = all_data[:, :, :-1]
    labels = all_data[:, :, -1]

    sg_indx = nt.config["core"]["data_types"]["signal"]

    images = np.copy(data[sg_indx])
    images = images.reshape(images.shape[0], -1)

    high_current_images = np.max(images, axis=1)
    high_current_ids = np.where(high_current_images > 1)[0]

    # print(len(high_current_ids))

    for exid in high_current_ids:
        # print(np.max(data[sg_indx, exid]))
        sig = data[sg_indx, exid]
        sig = (sig - np.min(sig)) / (np.max(sig) - np.min(sig))
        sig = sig * 0.3  # assume it's dots and highest current is not max current
        data[sg_indx, exid] = sig

        freq_spect = fp.fft2(sig.reshape(50, 50))
        freq_spect = np.abs(fp.fftshift(freq_spect))

        grad = generic_gradient_magnitude(sig.reshape(50, 50), sobel)

        index = nt.config["core"]["data_types"]["frequencies"]
        data[index, exid, :] = freq_spect.flatten()

        index = nt.config["core"]["data_types"]["gradient"]
        data[index, exid, :] = grad.flatten()

    n = list(data.shape)
    n[-1] += 1

    data_w_labels = np.zeros(n)
    data_w_labels[:, :, -1] = labels
    data_w_labels[:, :, :-1] = data

    path = os.path.join(db_folder, filename)
    np.save(path, data_w_labels)


"""This module provides EIGENVAL."""

from __future__ import division, print_function

import csv
import sys
from pathlib import Path
from logging import DEBUG, INFO, Formatter, StreamHandler, getLogger

import numpy as np
from typing import Dict, Optional, Sequence, Tuple, List, IO, Union
from vaspy.tools import open_by_suffix

try:
    import matplotlib.pyplot as plt
except ImportError:
    sys.stderr.write("Install matplotlib, or you cannot use methods relating to draw\n")

# logger
LOGLEVEL = INFO
logger = getLogger(__name__)
fmt = "%(asctime)s %(levelname)s %(name)s :%(message)s"
formatter = Formatter(fmt)
handler = StreamHandler()
handler.setLevel(LOGLEVEL)
logger.setLevel(LOGLEVEL)
handler.setFormatter(formatter)
logger.addHandler(handler)
logger.propagate = False

class EnergyBand(object):
    """Simple band structure object for analyzing by using ipython.

    Class for band structure

    Attributes
    ----------
    kvecs: numpy.ndarray
        kvectors
    kdistances: numpy.ndarray
        kdisance
    numk: int
        number of kpoints
    nbands: int
        number of bands
    nsping: int
        spin character
    energies: numpy.ndarray
        energies[spin_i, k_i, band_i], where spin_i, k_i, and band_i are spin-,
        k- and band-index, respectively.
    label: dict
        used as a label (data 'title' such as '#k', 'Energy') in str format

    Parameters
    ----------
    kvecs: numpy.ndarray
            1D array data of k-vectors.
    energies: numpy.ndarray
            1D array data of energies
    nspin: int
            number of spin: '1' means No-spin.  '2' means collinear spin,
            '4' means noncollinear spin.
            In this class does not distinguish non-collinear spin
            and No-spin.  (default is 1)

    """

    def __init__(
        self,
        kvecs: Sequence[float] = (),
        energies: Sequence[float] = (),
        nspin: int = 1,
    ) -> None:
        """Initialize."""
        self.kvecs: np.ndarray = np.array(kvecs)
        self.numk: int = len(self.kvecs)
        self.label: Dict["str", List["str"]] = {}
        try:
            self.nbands: int = len(energies) // len(kvecs)
        except ZeroDivisionError:
            self.nbands = 0
        self.energies: np.ndarray = np.array(energies)
        self.nspin = nspin
        if self.nspin == 1:  # standard
            self.label["spin"] = [""]
            self.label["energy"] = ["Energy"]
        elif self.nspin == 2:  # spin-polarized
            self.label["energy"] = ["Energy_up", "Energy_down"]
            self.label["spin"] = ["_up", "_down"]
        elif self.nspin == 4:  # non-collinear
            self.label["energy"] = ["Energy"]
            self.label["spin"] = ["_mT", "_mX", "_mY", "_mZ"]
        self.label["k"] = ["#k"]

    @property
    def kdistances(self) -> np.ndarray:
        """Return kdistances."""
        return np.cumsum(
            np.linalg.norm(
                np.concatenate((np.array([[0, 0, 0]]), np.diff(self.kvecs, axis=0))),
                axis=1,
            )
        )

    def fermi_correction(self, fermi: float) -> None:
        """Correct the Fermi level.

        Parameters
        ----------
        fermi: float
            value of the Fermi level.

        """
        self.energies -= fermi

    def make_label(self, *keys: str) -> List[str]:
        """Return array the used for label for CSV-like data.

        Parameters
        ----------
        keys: tuple
            key tuple used for label
s
        """
        label_list = []
        for key in keys:
            for tmp in self.label[key]:
                label_list.append(tmp)
        return label_list

    def to_3dlist(self) -> List[List[List[float]]]:
        """Return 3D mentional list.

        list[band_i, [k_i, energy, (energy_down)]]

        This list format would be useful for str output

        """
        bandstructure = []
        for energies in self.energies.T.tolist():
            band = []
            for k, energy in zip(self.kdistances[:, np.newaxis].tolist(), energies):
                k.extend(energy)
                band.append(k)
            bandstructure.append(band)
        return bandstructure

    def to_csv(self, csv_file: str, blankline: bool = True) -> None:
        """Write data to csv file.

        Parameters
        ------------
        csv_file: str
            filename for output
        label_str: str
            string for label (put it on the first line)
        blankline: boolean
            It True (default), the blank line is inserted between band data

        """
        label_str: str = "\t".join(self.make_label("k", "energy")) + "\n"
        with open(csv_file, "w") as fhandle:
            fhandle.writelines(label_str)
            writer = csv.writer(fhandle, delimiter="\t")
            for band_i in self.to_3dlist():
                writer.writerows(band_i)
                if blankline:
                    fhandle.writelines("\n")

    def __str__(self) -> str:
        """Return the str object.

        Returns
        --------
        str
            a string represntation of EnergyBand.
            **Useful for gnuplot and Igor**.

        """
        labels = self.make_label("k", "energy")
        output = labels[0]
        for label in labels[1:]:
            output += "\t" + label
        output += "\n"
        list3d = self.to_3dlist()
        for band_i in list3d:
            for line in band_i:
                output += "{0:.8e}".format(line[0])
                for energy in line[1:]:
                    output += "\t{0:.8e}".format(energy)
                output += "\n"
            output += "\n"
        return output

    def figure(self, color: str = "blue", spin_i: int = 0) -> plt.Axes:
        """Return Axes object of the energy band.

        Parameters
        -----------
        color: str, optional (default is 'blue')
            color of the band line

        spin_i: spin_index
            default is 0

        Returns
        ---------
        matplotlib.pyplot.Axes

        Example
        --------
        Here is a typical code::

            fig = plt.figure()
            ax = band.figure(color='blue')
            ax.set_ylabel('Energy  ( eV )')
            ax.set_ylim(-5, 5)
            ax.set_xlim(0, 4)
            plt.show()

        """
        [
            plt.plot(self.kdistances, self.energies[spin_i, :, band_i], color=color)
            for band_i in range(self.energies.shape[2])
        ]
        return plt.gca()

    def show(
        self, yrange: Optional[Tuple[float, float]] = None, spin_i: int = 0
    ) -> None:  # How to set default value?
        """Draw band structure by using maptlotlib.

        For 'just seeing' use.

        Parameters
        ----------
        yrange: tuple, optional  (default: all range)
            Minimum and maximum value of the y-axis.
            If not specified, use the matplotlib default value.

        spin_i: int  (default is 0 for no spin or 'up' spin)
            Spin index. For spin-polarized collinear band

        """
        for band_i in range(self.energies.shape[2]):
            plt.plot(self.kdistances, self.energies[spin_i, :, band_i], color="blue")
        if yrange is not None:
            plt.ylim([yrange[0], yrange[1]])
        plt.xlim([self.kdistances[0], self.kdistances[-1]])
        plt.ylabel(self.label["energy"][spin_i] + " (eV)")
        plt.show()

    def to_physical_kvector(
        self,
        recvec: np.ndarray = np.array(
            ((1.0, 0.0, 0.0), (0.0, 1.0, 0.0), (0.0, 0.0, 1.0)),
        ),
    ) -> None:
        """Change kvec unit to inverse AA.

        Parameters
        -----------
        recvec: array, numpy.ndarray, optional (default is the unit vector)
            reciprocal vector

        Notes
        -----
            Don't forget that the reciprocal vector used
            in VASP needs 2PI to match  the conventional
            unit of the wavevector.

        """
        logger.debug("recvec: {}".format(recvec))
        logger.debug("self.kvecs: {}".format(self.kvecs))
        recvec = np.array(recvec)
        self.kvecs = np.array([recvec.dot(kvecs) for kvecs in self.kvecs])

class EIGENVAL(EnergyBand):
    """Class for storing the data of EIGENVAL file.

    Parameters
    -----------
    filename: str, Path
        File name of 'EIGENVAL'

    Attributes
    ----------
    natom: int
        Number of atoms

    """

    def __init__(self, filename: Union[str, Path, None] = None) -> None:
        """Initialize."""
        super(EIGENVAL, self).__init__()
        self.natom = 0
        #
        if filename:
            self.load_file(open_by_suffix(str(filename)))

    def __getitem__(self, item: int) -> Tuple[List[float], List[List[float]]]:
        """

        Parameters
        ----------
        item: int
            index of k-vector

        Returns
        -------
        Tuple of list of float and list of float
        """
        energies: List[List[List[float]]] = self.energies.transpose(1, 2, 0).tolist()
        kvec: List[List[float]] = self.kvecs.tolist()
        return list(zip(kvec, energies))[item]

    def __len__(self) -> int:
        """Return numk as the result of len()"""
        return self.numk

    def load_file(self, thefile: IO[str]) -> None:
        """Parse EIGENVAL."""
        self.natom, _, _, self.nspin = [int(i) for i in next(thefile).split()]
        if self.nspin == 2:
            self.label["energy"] = ["Energy_up", "Energy_down"]
        else:
            self.label["energy"] = ["Energy"]
        next(thefile)
        next(thefile)
        next(thefile)
        next(thefile)
        _, self.numk, self.nbands = [int(i) for i in next(thefile).split()]
        self.kvecs = []
        self.energies = []
        for _ in range(self.numk):
            # the first line in the sigleset begins with the blank
            next(thefile)
            self.kvecs.append([float(i) for i in next(thefile).split()[0:3]])
            for _ in range(self.nbands):
                self.energies.append(
                    [float(i) for i in next(thefile).split()[1 : self.nspin + 1]]
                )
        self.kvecs = np.array(self.kvecs)
        self.energies = np.array(self.energies).T.reshape(
            self.nspin, self.numk, self.nbands
        )
        thefile.close()


import os
import numpy as np
import pandas as pd
from .chebyshevUtils import chebeval

__all__ = ['ChebyValues']

class ChebyValues(object):
    """Calculates positions, velocities, deltas, vmags and elongations,
    given a series of coefficients generated by ChebyFits.
    """
    def __init__(self):
        self.coeffs = {}
        self.coeffKeys = ['objId', 'tStart', 'tEnd', 'ra', 'dec', 'geo_dist', 'vmag', 'elongation']
        self.ephemerisKeys = ['ra', 'dradt', 'dec', 'ddecdt', 'geo_dist', 'vmag', 'elongation']

    def setCoefficients(self, chebyFits):
        """Set coefficients using a ChebyFits object.
        (which contains a dictionary of objId, tStart, tEnd, ra, dec, delta, vmag, and elongation lists).

        Parameters
        ----------
        chebyFits : chebyFits
            ChebyFits object, with attribute 'coeffs' - a dictionary of lists of coefficients.
        """
        self.coeffs = chebyFits.coeffs
        # Convert list of coefficients into numpy arrays.
        for k in self.coeffs:
            self.coeffs[k] = np.array(self.coeffs[k])
        # Check that expected values were received.
        missing_keys = set(self.coeffKeys) - set(self.coeffs)
        if len(missing_keys) > 0:
            raise ValueError("Expected to find key(s) %s in coefficients." %  ' '.join(list[missing_keys]))
        self.coeffs['meanRA'] = self.coeffs['ra'].swapaxes(0, 1)[0]
        self.coeffs['meanDec'] = self.coeffs['dec'].swapaxes(0, 1)[0]

    def readCoefficients(self, chebyFitsFile):
        """Read coefficients from output file written by ChebyFits.

        Parameters
        ----------
        chebyFitsFile : str
            The filename of the coefficients file.
        """
        if not os.path.isfile(chebyFitsFile):
            raise IOError('Could not find chebyFitsFile at %s' % (chebyFitsFile))
        # Read the coefficients file.
        coeffs = pd.read_table(chebyFitsFile, delim_whitespace=True)
        # The header line provides information on the number of coefficients for each parameter.
        datacols = coeffs.columns.values
        cols = {}
        coeff_cols = ['ra', 'dec', 'geo_dist', 'vmag', 'elongation']
        for k in coeff_cols:
            cols[k] = [x for x in datacols if x.startswith(k)]
        # Translate dataframe to dictionary of numpy arrays
        # while consolidating RA/Dec/Delta/Vmag/Elongation coeffs.
        self.coeffs['objId'] = coeffs.objId.values
        self.coeffs['tStart'] = coeffs.tStart.values
        self.coeffs['tEnd'] = coeffs.tEnd.values
        for k in coeff_cols:
            self.coeffs[k] = np.empty([len(cols[k]), len(coeffs)], float)
            for i in range(len(cols[k])):
                self.coeffs[k][i] = coeffs['%s_%d' % (k, i)].values
        # Add the mean RA and Dec columns (before swapping the coefficients axes).
        self.coeffs['meanRA'] = self.coeffs['ra'][0]
        self.coeffs['meanDec'] = self.coeffs['dec'][0]
        # Swap the coefficient axes so that they are [segment, coeff].
        for k in coeff_cols:
            self.coeffs[k] = self.coeffs[k].swapaxes(0, 1)

    def _evalSegment(self, segmentIdx, times, subsetSegments=None, mask=True):
        """Evaluate the ra/dec/delta/vmag/elongation values for a given segment at a series of times.

        Parameters
        ----------
        segmentIdx : int
            The index in (each of) self.coeffs for the segment.
            e.g. the first segment, for each object.
        times : np.ndarray
            The times at which to evaluate the segment.
        subsetSegments : numpy.ndarray, optional
            Optionally specify a subset of the total segment indexes.
            This lets you pick out particular objIds.
        mask : bool, optional
            If True, returns NaNs for values outside the range of times in the segment.
            If False, extrapolates segment for times outside the segment time range.

        Returns
        -------
        dict
           Dictionary of RA, Dec, delta, vmag, and elongation values for the segment indicated,
           at the time indicated.
        """
        if subsetSegments is None:
            subsetSegments = np.ones(len(self.coeffs['objId']), dtype=bool)
        tStart = self.coeffs['tStart'][subsetSegments][segmentIdx]
        tEnd = self.coeffs['tEnd'][subsetSegments][segmentIdx]
        tScaled = times - tStart
        tInterval = np.array([tStart, tEnd]) - tStart
        # Evaluate RA/Dec/Delta/Vmag/elongation.
        ephemeris = {}
        ephemeris['ra'], ephemeris['dradt'] = chebeval(tScaled,
                                                       self.coeffs['ra'][subsetSegments][segmentIdx],
                                                       interval=tInterval, doVelocity=True, mask=mask)
        ephemeris['dec'], ephemeris['ddecdt'] = chebeval(tScaled,
                                                         self.coeffs['dec'][subsetSegments][segmentIdx],
                                                         interval=tInterval, doVelocity=True, mask=mask)
        ephemeris['dradt'] = ephemeris['dradt'] * np.cos(np.radians(ephemeris['dec']))
        for k in ('geo_dist', 'vmag', 'elongation'):
            ephemeris[k], _ = chebeval(tScaled, self.coeffs[k][subsetSegments][segmentIdx],
                                       interval=tInterval, doVelocity=False, mask=mask)
        return ephemeris

    def getEphemerides(self, times, objIds=None, extrapolate=False):
        """Find the ephemeris information for 'objIds' at 'time'.

        Implicit in how this is currently written is that the segments are all expected to cover the
        same start/end time range across all objects.
        They do not have to have the same segment length for all objects.

        Parameters
        ----------
        times : float or np.ndarray
            The time to calculate ephemeris positions.
        objIds : numpy.ndarray, optional
            The object ids for which to generate ephemerides. If None, then just uses all objects.
        extrapolate : bool
            If True, extrapolate beyond ends of segments if time outside of segment range.
            If False, return ValueError if time is beyond range of segments.

        Returns
        -------
        numpy.ndarray
            The ephemeris positions for all objects.
            Note that these may not be sorted in the same order as objIds.
        """
        if isinstance(times, float) or isinstance(times, int):
            times = np.array([times], float)
        ntimes = len(times)
        ephemerides = {}
        # Find subset of segments which match objId, if specified.
        if objIds is None:
            objMatch = np.ones(len(self.coeffs['objId']), dtype=bool)
            ephemerides['objId'] = np.unique(self.coeffs['objId'])
        else:
            if isinstance(objIds, str) or isinstance(objIds, int):
                objIds = np.array([objIds])
            objMatch = np.in1d(self.coeffs['objId'], objIds)
            ephemerides['objId'] = objIds
        # Now find ephemeris values.
        ephemerides['time'] = np.zeros((len(ephemerides['objId']), ntimes), float) + times
        for k in self.ephemerisKeys:
            ephemerides[k] = np.zeros((len(ephemerides['objId']), ntimes), float)
        for it, t in enumerate(times):
            # Find subset of segments which contain the appropriate time.
            # Look for simplest subset first.
            segments = np.where((self.coeffs['tStart'][objMatch] <= t) &
                                (self.coeffs['tEnd'][objMatch] > t))[0]
            if len(segments) == 0:
                segStart = self.coeffs['tStart'][objMatch].min()
                segEnd = self.coeffs['tEnd'][objMatch].max()
                if (segStart > t or segEnd < t):
                    if not extrapolate:
                        for k in self.ephemerisKeys:
                            ephemerides[k][:,it] = np.nan
                    else:
                        # Find the segments to use to extrapolate the times.
                        if segStart > t:
                            segments = np.where(self.coeffs['tStart'][objMatch] == segStart)[0]
                        if segEnd < t:
                            segments = np.where(self.coeffs['tEnd'][objMatch] == segEnd)[0]
                elif segEnd == t:
                    # Not extrapolating, but outside the simple match case above.
                    segments = np.where(self.coeffs['tEnd'][objMatch] == segEnd)[0]
            for i, segmentIdx in enumerate(segments):
                ephemeris = self._evalSegment(segmentIdx, t, objMatch, mask=False)
                for k in self.ephemerisKeys:
                    ephemerides[k][i][it] = ephemeris[k]
                ephemerides['objId'][i] = self.coeffs['objId'][objMatch][segmentIdx]
        if objIds is not None:
            if set(ephemerides['objId']) != set(objIds):
                raise ValueError('Did not find expected match between objIds provided and ephemeride objIds.')
        return ephemerides


import numpy as np
import re
from pymatgen.analysis.local_env import VoronoiNN
from pymatgen.core.structure import Structure
import math
import os
import csv
import argparse
import sys
import time

parser = argparse.ArgumentParser(description='Orbital Field Matrix')
parser.add_argument('root', help='path to the directory of CIF files.')
parser.add_argument('--partial_csv', type=str, default=None, help='the path to the csv file with specified atoms')
args = parser.parse_args(sys.argv[1:])

cif_path = args.root
partial_csv = args.partial_csv

start_all = time.time()
# ----------------------------- Pre-defined dictionary -------------------------------------------

elements = {'H': ['1s2'], 'Li': ['[He] 1s2'], 'Be': ['[He] 2s2'], 'B': ['[He] 2s2 2p1'], 'N': ['[He] 2s2 2p3'],
            'O': ['[He] 2s2 2p4'],
            'C': ['[He] 2s2 2p2'], 'I': ['[Kr] 4d10 5s2 5p5'],
            'F': ['[He] 2s2 2p5'], 'Na': ['[Ne] 3s1'], 'Mg': ['[Ne] 3s2'], 'Al': ['[Ne] 3s2 3p1'],
            'Si': ['[Ne] 3s2 3p2'],
            'P': ['[Ne] 3s2 3p3'], 'S': ['[Ne] 3s2 3p4'], 'Cl': ['[Ne] 3s2 3p5'], 'K': ['[Ar] 4s1'],
            'Ca': ['[Ar] 4s2'], 'Sc': ['[Ar] 3d1 4s2'],
            'Ti': ['[Ar] 3d2 4s2'], 'V': ['[Ar] 3d3 4s2'], 'Cr': ['[Ar] 3d5 4s1'], 'Mn': ['[Ar] 3d5 4s2'],
            'Fe': ['[Ar] 3d6 4s2'], 'Co': ['[Ar] 3d7 4s2'], 'Ni': ['[Ar] 3d8 4s2'], 'Cu': ['[Ar] 3d10 4s1'],
            'Zn': ['[Ar] 3d10 4s2'],
            'Ga': ['[Ar] 3d10 4s2 4p2'], 'Ge': ['[Ar] 3d10 4s2 4p2'], 'As': ['[Ar] 3d10 4s2 4p3'],
            'Se': ['[Ar] 3d10 4s2 4p4'], 'Br': ['[Ar] 3d10 4s2 4p5'], 'Rb': ['[Kr] 5s1'],
            'Sr': ['[Kr] 5s2'], 'Y': ['[Kr] 4d1 5s2'], 'Zr': ['[Kr] 4d2 5s2'], 'Nb': ['[Kr] 4d4 5s1'],
            'Mo': ['[Kr] 4d5 5s1'],
            'Ru': ['[Kr] 4d7 5s1'], 'Rh': ['[Kr] 4d8 5s1'], 'Pd': ['[Kr] 4d10'], 'Ag': ['[Kr] 4d10 5s1'],
            'Cd': ['[Kr] 4d10 5s2'],
            'In': ['[Kr] 4d10 5s2 5p1'], 'Sn': ['[Kr] 4d10 5s2 5p2'], 'Sb': ['[Kr] 4d10 5s2 5p3'],
            'Te': ['[Kr] 4d10 5s2 5p4'], 'Cs': ['[Xe] 6s1'], 'Ba': ['[Xe] 6s2'],
            'La': ['[Xe] 5d1 6s2'], 'Ce': ['[Xe] 4f1 5d1 6s2'], 'Hf': ['[Xe] 4f14 5d2 6s2'],
            'Ta': ['[Xe] 4f14 5d3 6s2'],
            'W': ['[Xe] 4f14 5d5 6s1'], 'Re': ['[Xe] 4f14 5d5 6s2'], 'Os': ['[Xe] 4f14 5d6 6s2'],
            'Ir': ['[Xe] 4f14 5d7 6s2'], 'Pt': ['[Xe] 4f14 5d10'], 'Au': ['[Xe] 4f14 5d10 6s1'],
            'Hg': ['[Xe] 4f14 5d10 6s2'],
            'Tl': ['[Xe] 4f14 5d10 6s2 6p2'], 'Pb': ['[Xe] 4f14 5d10 6s2 6p2'],
            'Bi': ['[Xe] 4f14 5d10 6s2 6p3'],
            'Tc': ['[Kr] 4d5 5s2'], 'Fr': ['[Rn]7s1'], 'Ra': ['[Rn]7s2'], 'Pr': ['[Xe]4f3 6s2'],
            'Nd': ['[Xe] 4f4 6s2'], 'Pm': ['[Xe] 4f5 6s2'], 'Sm': ['[Xe] 4f6 6s2'],
            'Eu': ['[Xe] 4f7 6s2'], 'Gd': ['[Xe] 4f7 5d1 6s2'], 'Tb': ['[Xe] 4f9 6s2'],
            'Dy': ['[Xe] 4f10 6s2'], 'Ho': ['[Xe] 4f11 6s2'], 'Er': ['[Xe] 4f12 6s2'],
            'Tm': ['[Xe] 4f13 6s2'], 'Yb': ['[Xe] 4f14 6s2'], 'Lu': ['[Xe] 4f14 5d1 6s2'],
            'Po': ['[Xe] 4f14 5d10 6s2 6p4'], 'At': ['[Xe] 4f14 5d10 6s2 6p5'],
            'Ac': ['[Rn] 6d1 7s2'], 'Th': ['[Rn] 6d2 7s2'], 'Pa': ['[Rn] 5f2 6d1 7s2'],
            'U': ['[Rn] 5f3 6d1 7s2'], 'Np': ['[Rn] 5f4 6d1 7s2'], 'Pu': ['[Rn] 5f6 7s2'],
            'Am': ['[Rn] 5f7 7s2'], 'Cm': ['[Rn] 5f7 6d1 7s2'], 'Bk': ['[Rn] 5f9 7s2'],
            'Cf': ['[Rn] 5f10 7s2'], 'Es': ['[Rn] 5f11 7s2'], 'Fm': ['[Rn] 5f12 7s2'],
            'Md': ['[Rn] 5f13 7s2'], 'No': ['[Rn] 5f14 7s2'], 'Lr': ['[Rn] 5f14 6d1 7s2'],
            'Rf': ['[Rn] 5f14 6d2 7s2'], 'Db': ['[Rn] 5f14 6d3 7s2'],
            'Sg': ['[Rn] 5f14 6d4 7s2'], 'Bh': ['[Rn] 5f14 6d5 7s2'],
            'Hs': ['[Rn] 5f14 6d6 7s2'], 'Mt': ['[Rn] 5f14 6d7 7s2'], 'Xe': ['[Kr] 4d10 5s2 5p6'],
            'He': ['1s2'], 'Kr': ['[Ar] 3d10 4s2 4p6'], 'Ar': ['[Ne] 3s2 3p6'], 'Ne': ['[He] 2s2 2p6']}

orbitals = {"s1": 0, "s2": 1, "p1": 2, "p2": 3, "p3": 4, "p4": 5, "p5": 6, "p6": 7, "d1": 8, "d2": 9, "d3": 10,
            "d4": 11,
            "d5": 12, "d6": 13, "d7": 14, "d8": 15, "d9": 16, "d10": 17, "f1": 18, "f2": 19, "f3": 20, "f4": 21,
            "f5": 22, "f6": 23, "f7": 24, "f8": 25, "f9": 26, "f10": 27, "f11": 28, "f12": 29, "f13": 30,
            "f14": 31}

# ------------------------- hvs ----------------------------------------------------
# hvs define a dictionary , which map a element to a 32 vector representation
# according to its electronic configurations

hvs = {}

for key in elements.keys():
    element = key
    hv = np.zeros(shape=(32, 1))
    s = elements[key][0]
    sp = (re.split('(\s+)', s))
    if key == "H":
        hv[0] = 1
    if key != "H":
        for j in range(1, len(sp)):
            if sp[j] != ' ':
                n = sp[j][:1]
                orb = sp[j][1:]
                hv[orbitals[orb]] = 1
    hvs[element] = hv

# --------------------------- pre-defined functions ----------------------------------

def make_hot_for_atom_i(crystal, i, hvs):
    EP = str(crystal[i].specie)
    # nan_to_num: Replace nan with zero and inf with finite numbers.
    HV_P = np.nan_to_num(hvs[EP])
    # reshape from (32,1) to (1,32)
    AA = HV_P.reshape((HV_P.shape[1], 32))
    A = np.array(AA)
    # get the Voronoi nearest neighbours(VNN) of atom indexed by i
    b = VoronoiNN().get_nn_info(crystal, i)
    angles = []
    # store the solid angles between central atom i with VNN's in angles
    for nb in b:
        angle_K = nb['poly_info']['solid_angle']
        angles.append(angle_K)
    max_angle = max(angles)
    X_P = np.zeros(shape=(32, 32))
    tmp_X = []
    for nb in b:
        # check VNN b's specie type
        EK = str(nb['site'].specie)
        # check the solid angle between b and central atom
        angle_K = nb['poly_info']['solid_angle']
        index_K = nb['site_index']
        # calculate the distance square between b and central atom
        r_pk = ((calculateDistance(nb['site'].coords, crystal[i].coords)) * (
            calculateDistance(nb['site'].coords, crystal[i].coords)))
        # map b to the vector representation and reshape
        HV_K = hvs[EK]
        HV_K = HV_K.reshape((HV_K.shape[1], 32))
        # weight coeficients of b-central atom pair
        coef_K = (angle_K / max_angle) * ((1 / ((r_pk) ** 2)))
        HV_K_new = np.nan_to_num(coef_K * HV_K)
        # matrix product between b and central atom ---> (32, 32)
        X_PT = np.matmul(HV_P, HV_K_new)
        tmp_X.append(X_PT)
    X0 = np.zeros(shape=(32, 32))
    # el stands for one matrix(of VNN-central_atom pair)
    # weighted sum over all pairs
    for el in tmp_X:
        X0 = [[sum(x) for x in zip(el[i], X0[i])] for i in range(len(el))]
    # the size of X0 is (32, 33),not flattened
    # the first column is central atom
    X0 = np.concatenate((A.T, X0), axis=1)
    X0 = np.asarray(X0)
    # flatten arranged by column
    X0 = X0.flatten(order='F')
    return X0

def calculateDistance(a, b):  # Atom-wise OFM
    dist = math.sqrt((a[0] - b[0]) ** 2 + (a[1] - b[1]) ** 2 + (a[2] - b[2]) ** 2)
    return dist

# ---------------------- Loading id-data(cif-id) of magnetic material-----------------------------
id_prop_file = os.path.join(cif_path, 'id_prop.csv')
with open(id_prop_file) as f:
    reader = csv.reader(f)
    id_prop_data = [row for row in reader]

id_data = []
for i in range(len(id_prop_data)):
    id_data.append(id_prop_data[i][0])

if args.partial_csv is not None:
    with open(args.partial_csv) as f:
        reader = csv.reader(f)
        magmom_data = [row for row in reader]
    id = []
    for i in range(len(magmom_data)):
        id.append(magmom_data[i][0])

# ----------------------- Consruct OFM representation for given material ao given position -----------------------
# store in a dictionary called all_atom_embedding
all_atom_embedding = {}
start = time.time()

for ids in id_data:
    crystal = Structure.from_file(os.path.join(cif_path, ids + '.cif'))
    for idx in range(len(crystal)):
        atom_name = ids + '_' + str(crystal[idx].specie) + '_' + str(crystal[idx].specie.number) + '_' + str(idx + 1)
        if args.partial_csv is None:
            all_atom_embedding[atom_name] = make_hot_for_atom_i(crystal, idx, hvs)
        else:
            if atom_name in id:
                all_atom_embedding[atom_name] = make_hot_for_atom_i(crystal, idx, hvs)

print("Spend ", time.time() - start, ' s to store OFM rep of atoms in a dictionary all_atom_embedding')
print("*********************************************************************")
# --------------------------- Save as .npy file ----------------------------------

np.save('OFM.npy', all_atom_embedding)
print("Over !!!!! This script takes: ", time.time() - start_all, ' s')


#!/usr/bin/env python
# coding=utf-8
"""
@author: Jiawei Wu
@create time: 2019-11-17 11:23
@edit time: 2020-05-10 22:17
@FilePath: /vvlab/agents/DQN_base.py
@desc: 创建DQN对象
"""
import torch
import torch.nn as nn
import numpy as np
from ..utils import ReplayBuffer
CUDA = torch.cuda.is_available()

class DQNBase(object):
    """The base class for DQN."""

    def __init__(self, n_states, n_actions, learning_rate=0.001,
                 discount_rate=0.0, card_no=0, **kwargs):
        """Initialize two networks and experience playback pool of DQN.

        Args:
          n_states: Number of states.
          n_actions: Number of actions.
          learning_rate: Decide how much error this time is to be learned.
          discount_rate: Attenuation value for future reward.
          card_no: Designated training card number.
          **kwargs: Incoming parameters.
        """
        # super parameters of DQN
        self.gamma = discount_rate  # future discount rate

        self.epsilon = 0.6
        self.epsilon_min = 0.001
        self.epsilon_decay = 0.999
        self.eval_every = 10
        self.card_no = card_no
        # bulid the network
        self.n_states, self.n_actions = n_states, n_actions
        self._build_net()

        self.buff_size, self.buff_thres, self.batch_size = 50000, 1000, 256
        if 'buff_size' in kwargs:
            self.buff_size = kwargs['buff_size']
        if 'buff_thres' in kwargs:
            self.buff_thres = kwargs['buff_thres']
        if 'batch_size' in kwargs:
            self.batch_size = kwargs['batch_size']
        self.replay_buff = ReplayBuffer(n_states, 1, buff_size=self.buff_size,
                                        buff_thres=self.buff_thres,
                                        batch_size=self.batch_size,
                                        card_no=self.card_no)
        # define optimizer and loss function
        self.optimizer = torch.optim.Adam(
            self.eval_net.parameters(), lr=learning_rate)
        self.loss_func = nn.MSELoss()
        # record the number of steps for synchronization parameters
        self.eval_step = 0
        if CUDA:
            self.cuda()

    def _build_net(self):
        """Bulid the network.

        Raises:
          TypeError:Network build no implementation.
        """
        raise TypeError("Network build no implementation")

    def get_action(self, state):
        """Get the action at this moment.

        Args:
          state:State at this moment.

        Return:
          The action output.
        """
        # epsilon update
        self.epsilon = self.epsilon * \
            self.epsilon_decay \
            if self.epsilon > self.epsilon_min else self.epsilon
        # convert row vector to column vector
        # (1 x n_states -> n_states x 1 x 1)
        if np.random.rand() < self.epsilon:
            # random
            action_size = state.shape[0]
            return np.random.randint(0, self.n_actions, (1, action_size))
        else:
            # greedy
            state = torch.unsqueeze(torch.FloatTensor(state), 0)
            action_values = self.eval_net.forward(state).cpu()
            return action_values.data.numpy().argmax(axis=2)

    def get_raw_out(self, state):
        """Get the original actions.

        Args:
          state:State at this moment.

        Returns:
          Action values before selected by argmax.
        """
        state = torch.unsqueeze(torch.FloatTensor(state), 0)
        action_values = self.eval_net.forward(state)
        print(action_values)
        return action_values

    def add_step(self, cur_state, action, reward, done, next_state):
        """Add a record to the experience replay pool.

        Args:
          cur_state:State at this moment.
          action:Action output at this moment.
          reward:Reward after taking the action.
          done:A sign to indicate whether training is stopped.
          next_state:State at next moment.
        """
        self.replay_buff.add_step(cur_state, action, reward, done, next_state)

    def learn(self):
        """The learning process of DQN.

        Returns:
          The loss during learning process when the batch is not 0.
        """
        batch = self.replay_buff.get_batch_splited_tensor(CUDA)
        if batch is None:
            return None
        # copy parameter
        if self.eval_step % self.eval_every == 0:
            self.target_net.load_state_dict(self.eval_net.state_dict())
        # update training steps
        self.eval_step += 1
        # split batch
        batch_cur_states, batch_actions, \
            batch_rewards, batch_dones, batch_next_states = batch
        # calculation error
        q_eval = self.eval_net(batch_cur_states)
        q_eval = q_eval.gather(1, batch_actions.long())  # shape (batch, 1)
        # detach from graph, don't backpropagate
        q_next = self.target_net(batch_next_states).detach()
        # if done, the future is not considered
        q_target = batch_rewards + self.gamma * \
            (1 - batch_dones) * \
            q_next.max(1)[0].view(
                len(batch_next_states), 1)   # shape (batch, 1)
        loss = self.loss_func(q_eval, q_target)
        # update network
        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()
        return loss.detach().cpu().numpy()

    def cuda(self):
        """GPU operation using specified card."""
        self.eval_net.cuda(self.card_no)
        self.target_net.cuda(self.card_no)


from typing import Dict, Any

import torch
import math
import logging

import numpy as np

_logger = logging.getLogger(__name__)

class Scheduler:
    """ Parameter Scheduler Base Class
    A scheduler base class that can be used to schedule any optimizer parameter groups.
    Unlike the builtin PyTorch schedulers, this is intended to be consistently called
    * At the END of each epoch, before incrementing the epoch count, to calculate next epoch's value
    * At the END of each optimizer update, after incrementing the update count, to calculate next update's value
    The schedulers built on this should try to remain as stateless as possible (for simplicity).
    This family of schedulers is attempting to avoid the confusion of the meaning of 'last_epoch'
    and -1 values for special behaviour. All epoch and update counts must be tracked in the training
    code and explicitly passed in to the schedulers on the corresponding step or step_update call.
    Based on ideas from:
     * https://github.com/pytorch/fairseq/tree/master/fairseq/optim/lr_scheduler
     * https://github.com/allenai/allennlp/tree/master/allennlp/training/learning_rate_schedulers
    """

    def __init__(self,
                 optimizer: torch.optim.Optimizer,
                 param_group_field: str,
                 noise_range_t=None,
                 noise_type='normal',
                 noise_pct=0.67,
                 noise_std=1.0,
                 noise_seed=None,
                 initialize: bool = True) -> None:
        self.optimizer = optimizer
        self.param_group_field = param_group_field
        self._initial_param_group_field = f"initial_{param_group_field}"
        if initialize:
            for i, group in enumerate(self.optimizer.param_groups):
                if param_group_field not in group:
                    raise KeyError(f"{param_group_field} missing from param_groups[{i}]")
                group.setdefault(self._initial_param_group_field, group[param_group_field])
        else:
            for i, group in enumerate(self.optimizer.param_groups):
                if self._initial_param_group_field not in group:
                    raise KeyError(f"{self._initial_param_group_field} missing from param_groups[{i}]")
        self.base_values = [group[self._initial_param_group_field] for group in self.optimizer.param_groups]
        self.metric = None  # any point to having this for all?
        self.noise_range_t = noise_range_t
        self.noise_pct = noise_pct
        self.noise_type = noise_type
        self.noise_std = noise_std
        self.noise_seed = noise_seed if noise_seed is not None else 42
        self.update_groups(self.base_values)

    def state_dict(self) -> Dict[str, Any]:
        return {key: value for key, value in self.__dict__.items() if key != 'optimizer'}

    def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
        self.__dict__.update(state_dict)

    def get_epoch_values(self, epoch: int):
        return None

    def get_update_values(self, num_updates: int):
        return None

    def step(self, epoch: int, metric: float = None) -> None:
        self.metric = metric
        values = self.get_epoch_values(epoch)
        if values is not None:
            values = self._add_noise(values, epoch)
            self.update_groups(values)

    def step_update(self, num_updates: int, metric: float = None):
        self.metric = metric
        values = self.get_update_values(num_updates)
        if values is not None:
            values = self._add_noise(values, num_updates)
            self.update_groups(values)

    def update_groups(self, values):
        if not isinstance(values, (list, tuple)):
            values = [values] * len(self.optimizer.param_groups)
        for param_group, value in zip(self.optimizer.param_groups, values):
            param_group[self.param_group_field] = value

    def _add_noise(self, lrs, t):
        if self.noise_range_t is not None:
            if isinstance(self.noise_range_t, (list, tuple)):
                apply_noise = self.noise_range_t[0] <= t < self.noise_range_t[1]
            else:
                apply_noise = t >= self.noise_range_t
            if apply_noise:
                g = torch.Generator()
                g.manual_seed(self.noise_seed + t)
                if self.noise_type == 'normal':
                    while True:
                        # resample if noise out of percent limit, brute force but shouldn't spin much
                        noise = torch.randn(1, generator=g).item()
                        if abs(noise) < self.noise_pct:
                            break
                else:
                    noise = 2 * (torch.rand(1, generator=g).item() - 0.5) * self.noise_pct
                lrs = [v + v * noise for v in lrs]
        return lrs

class CosineLRScheduler(Scheduler):
    """
    Cosine decay with restarts.
    This is described in the paper https://arxiv.org/abs/1608.03983.
    Inspiration from
    https://github.com/allenai/allennlp/blob/master/allennlp/training/learning_rate_schedulers/cosine.py
    """

    def __init__(self,
                 optimizer: torch.optim.Optimizer,
                 t_initial: int,
                 t_mul: float = 1.,
                 lr_min: float = 0.,
                 decay_rate: float = 1.,
                 warmup_t=0,
                 warmup_lr_init=0,
                 warmup_prefix=False,
                 cycle_limit=0,
                 t_in_epochs=True,
                 noise_range_t=None,
                 noise_pct=0.67,
                 noise_std=1.0,
                 noise_seed=42,
                 initialize=True) -> None:
        super().__init__(
            optimizer, param_group_field="lr",
            noise_range_t=noise_range_t, noise_pct=noise_pct, noise_std=noise_std, noise_seed=noise_seed,
            initialize=initialize)

        assert t_initial > 0
        assert lr_min >= 0
        if t_initial == 1 and t_mul == 1 and decay_rate == 1:
            _logger.warning("Cosine annealing scheduler will have no effect on the learning "
                           "rate since t_initial = t_mul = eta_mul = 1.")
        self.t_initial = t_initial
        self.t_mul = t_mul
        self.lr_min = lr_min
        self.decay_rate = decay_rate
        self.cycle_limit = cycle_limit
        self.warmup_t = warmup_t
        self.warmup_lr_init = warmup_lr_init
        self.warmup_prefix = warmup_prefix
        self.t_in_epochs = t_in_epochs
        if self.warmup_t:
            self.warmup_steps = [(v - warmup_lr_init) / self.warmup_t for v in self.base_values]
            super().update_groups(self.warmup_lr_init)
        else:
            self.warmup_steps = [1 for _ in self.base_values]

    def _get_lr(self, t):
        if t < self.warmup_t:
            lrs = [self.warmup_lr_init + t * s for s in self.warmup_steps]
        else:
            if self.warmup_prefix:
                t = t - self.warmup_t

            if self.t_mul != 1:
                i = math.floor(math.log(1 - t / self.t_initial * (1 - self.t_mul), self.t_mul))
                t_i = self.t_mul ** i * self.t_initial
                t_curr = t - (1 - self.t_mul ** i) / (1 - self.t_mul) * self.t_initial
            else:
                i = t // self.t_initial
                t_i = self.t_initial
                t_curr = t - (self.t_initial * i)

            gamma = self.decay_rate ** i
            lr_min = self.lr_min * gamma
            lr_max_values = [v * gamma for v in self.base_values]

            if self.cycle_limit == 0 or (self.cycle_limit > 0 and i < self.cycle_limit):
                lrs = [
                    lr_min + 0.5 * (lr_max - lr_min) * (1 + math.cos(math.pi * t_curr / t_i)) for lr_max in lr_max_values
                ]
            else:
                lrs = [self.lr_min for _ in self.base_values]

        return lrs

    def get_epoch_values(self, epoch: int):
        if self.t_in_epochs:
            return self._get_lr(epoch)
        else:
            return None

    def get_update_values(self, num_updates: int):
        if not self.t_in_epochs:
            return self._get_lr(num_updates)
        else:
            return None

    def get_cycle_length(self, cycles=0):
        if not cycles:
            cycles = self.cycle_limit
        cycles = max(1, cycles)
        if self.t_mul == 1.0:
            return self.t_initial * cycles
        else:
            return int(math.floor(-self.t_initial * (self.t_mul ** cycles - 1) / (1 - self.t_mul)))

class TanhLRScheduler(Scheduler):
    """
    Hyberbolic-Tangent decay with restarts.
    This is described in the paper https://arxiv.org/abs/1806.01593
    """

    def __init__(self,
                 optimizer: torch.optim.Optimizer,
                 t_initial: int,
                 lb: float = -6.,
                 ub: float = 4.,
                 t_mul: float = 1.,
                 lr_min: float = 0.,
                 decay_rate: float = 1.,
                 warmup_t=0,
                 warmup_lr_init=0,
                 warmup_prefix=False,
                 cycle_limit=0,
                 t_in_epochs=True,
                 noise_range_t=None,
                 noise_pct=0.67,
                 noise_std=1.0,
                 noise_seed=42,
                 initialize=True) -> None:
        super().__init__(
            optimizer, param_group_field="lr",
            noise_range_t=noise_range_t, noise_pct=noise_pct, noise_std=noise_std, noise_seed=noise_seed,
            initialize=initialize)

        assert t_initial > 0
        assert lr_min >= 0
        assert lb < ub
        assert cycle_limit >= 0
        assert warmup_t >= 0
        assert warmup_lr_init >= 0
        self.lb = lb
        self.ub = ub
        self.t_initial = t_initial
        self.t_mul = t_mul
        self.lr_min = lr_min
        self.decay_rate = decay_rate
        self.cycle_limit = cycle_limit
        self.warmup_t = warmup_t
        self.warmup_lr_init = warmup_lr_init
        self.warmup_prefix = warmup_prefix
        self.t_in_epochs = t_in_epochs
        if self.warmup_t:
            t_v = self.base_values if self.warmup_prefix else self._get_lr(self.warmup_t)
            self.warmup_steps = [(v - warmup_lr_init) / self.warmup_t for v in t_v]
            super().update_groups(self.warmup_lr_init)
        else:
            self.warmup_steps = [1 for _ in self.base_values]

    def _get_lr(self, t):
        if t < self.warmup_t:
            lrs = [self.warmup_lr_init + t * s for s in self.warmup_steps]
        else:
            if self.warmup_prefix:
                t = t - self.warmup_t

            if self.t_mul != 1:
                i = math.floor(math.log(1 - t / self.t_initial * (1 - self.t_mul), self.t_mul))
                t_i = self.t_mul ** i * self.t_initial
                t_curr = t - (1 - self.t_mul ** i) / (1 - self.t_mul) * self.t_initial
            else:
                i = t // self.t_initial
                t_i = self.t_initial
                t_curr = t - (self.t_initial * i)

            if self.cycle_limit == 0 or (self.cycle_limit > 0 and i < self.cycle_limit):
                gamma = self.decay_rate ** i
                lr_min = self.lr_min * gamma
                lr_max_values = [v * gamma for v in self.base_values]

                tr = t_curr / t_i
                lrs = [
                    lr_min + 0.5 * (lr_max - lr_min) * (1 - math.tanh(self.lb * (1. - tr) + self.ub * tr))
                    for lr_max in lr_max_values
                ]
            else:
                lrs = [self.lr_min * (self.decay_rate ** self.cycle_limit) for _ in self.base_values]
        return lrs

    def get_epoch_values(self, epoch: int):
        if self.t_in_epochs:
            return self._get_lr(epoch)
        else:
            return None

    def get_update_values(self, num_updates: int):
        if not self.t_in_epochs:
            return self._get_lr(num_updates)
        else:
            return None

    def get_cycle_length(self, cycles=0):
        if not cycles:
            cycles = self.cycle_limit
        cycles = max(1, cycles)
        if self.t_mul == 1.0:
            return self.t_initial * cycles
        else:
            return int(math.floor(-self.t_initial * (self.t_mul ** cycles - 1) / (1 - self.t_mul)))

class StepLRScheduler(Scheduler):
    """
    """

    def __init__(self,
                 optimizer: torch.optim.Optimizer,
                 decay_t: float,
                 decay_rate: float = 1.,
                 warmup_t=0,
                 warmup_lr_init=0,
                 t_in_epochs=True,
                 noise_range_t=None,
                 noise_pct=0.67,
                 noise_std=1.0,
                 noise_seed=42,
                 initialize=True,
                 ) -> None:
        super().__init__(
            optimizer, param_group_field="lr",
            noise_range_t=noise_range_t, noise_pct=noise_pct, noise_std=noise_std, noise_seed=noise_seed,
            initialize=initialize)

        self.decay_t = decay_t
        self.decay_rate = decay_rate
        self.warmup_t = warmup_t
        self.warmup_lr_init = warmup_lr_init
        self.t_in_epochs = t_in_epochs
        if self.warmup_t:
            self.warmup_steps = [(v - warmup_lr_init) / self.warmup_t for v in self.base_values]
            super().update_groups(self.warmup_lr_init)
        else:
            self.warmup_steps = [1 for _ in self.base_values]

    def _get_lr(self, t):
        if t < self.warmup_t:
            lrs = [self.warmup_lr_init + t * s for s in self.warmup_steps]
        else:
            lrs = [v * (self.decay_rate ** (t // self.decay_t)) for v in self.base_values]
        return lrs

    def get_epoch_values(self, epoch: int):
        if self.t_in_epochs:
            return self._get_lr(epoch)
        else:
            return None

    def get_update_values(self, num_updates: int):
        if not self.t_in_epochs:
            return self._get_lr(num_updates)
        else:
            return None

class PlateauLRScheduler(Scheduler):
    """Decay the LR by a factor every time the validation loss plateaus."""

    def __init__(self,
                 optimizer,
                 decay_rate=0.1,
                 patience_t=10,
                 verbose=True,
                 threshold=1e-4,
                 cooldown_t=0,
                 warmup_t=0,
                 warmup_lr_init=0,
                 lr_min=0,
                 mode='max',
                 noise_range_t=None,
                 noise_type='normal',
                 noise_pct=0.67,
                 noise_std=1.0,
                 noise_seed=None,
                 initialize=True,
                 ):
        super().__init__(optimizer, 'lr', initialize=initialize)

        self.lr_scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
            self.optimizer,
            patience=patience_t,
            factor=decay_rate,
            verbose=verbose,
            threshold=threshold,
            cooldown=cooldown_t,
            mode=mode,
            min_lr=lr_min
        )

        self.noise_range = noise_range_t
        self.noise_pct = noise_pct
        self.noise_type = noise_type
        self.noise_std = noise_std
        self.noise_seed = noise_seed if noise_seed is not None else 42
        self.warmup_t = warmup_t
        self.warmup_lr_init = warmup_lr_init
        if self.warmup_t:
            self.warmup_steps = [(v - warmup_lr_init) / self.warmup_t for v in self.base_values]
            super().update_groups(self.warmup_lr_init)
        else:
            self.warmup_steps = [1 for _ in self.base_values]
        self.restore_lr = None

    def state_dict(self):
        return {
            'best': self.lr_scheduler.best,
            'last_epoch': self.lr_scheduler.last_epoch,
        }

    def load_state_dict(self, state_dict):
        self.lr_scheduler.best = state_dict['best']
        if 'last_epoch' in state_dict:
            self.lr_scheduler.last_epoch = state_dict['last_epoch']

    # override the base class step fn completely
    def step(self, epoch, metric=None):
        if epoch <= self.warmup_t:
            lrs = [self.warmup_lr_init + epoch * s for s in self.warmup_steps]
            super().update_groups(lrs)
        else:
            if self.restore_lr is not None:
                # restore actual LR from before our last noise perturbation before stepping base
                for i, param_group in enumerate(self.optimizer.param_groups):
                    param_group['lr'] = self.restore_lr[i]
                self.restore_lr = None

            self.lr_scheduler.step(metric, epoch)  # step the base scheduler

            if self.noise_range is not None:
                if isinstance(self.noise_range, (list, tuple)):
                    apply_noise = self.noise_range[0] <= epoch < self.noise_range[1]
                else:
                    apply_noise = epoch >= self.noise_range
                if apply_noise:
                    self._apply_noise(epoch)

    def _apply_noise(self, epoch):
        g = torch.Generator()
        g.manual_seed(self.noise_seed + epoch)
        if self.noise_type == 'normal':
            while True:
                # resample if noise out of percent limit, brute force but shouldn't spin much
                noise = torch.randn(1, generator=g).item()
                if abs(noise) < self.noise_pct:
                    break
        else:
            noise = 2 * (torch.rand(1, generator=g).item() - 0.5) * self.noise_pct

        # apply the noise on top of previous LR, cache the old value so we can restore for normal
        # stepping of base scheduler
        restore_lr = []
        for i, param_group in enumerate(self.optimizer.param_groups):
            old_lr = float(param_group['lr'])
            restore_lr.append(old_lr)
            new_lr = old_lr + old_lr * noise
            param_group['lr'] = new_lr
        self.restore_lr = restore_lr

def create_scheduler(args, optimizer):
    num_epochs = args.epochs

    if getattr(args, 'lr_noise', None) is not None:
        lr_noise = getattr(args, 'lr_noise')
        if isinstance(lr_noise, (list, tuple)):
            noise_range = [n * num_epochs for n in lr_noise]
            if len(noise_range) == 1:
                noise_range = noise_range[0]
        else:
            noise_range = lr_noise * num_epochs
    else:
        noise_range = None

    lr_scheduler = None
    if args.sched == 'cosine':
        lr_scheduler = CosineLRScheduler(
            optimizer,
            t_initial=num_epochs,
            t_mul=getattr(args, 'lr_cycle_mul', 1.),
            lr_min=args.min_lr,
            decay_rate=args.decay_rate,
            warmup_lr_init=args.warmup_lr,
            warmup_t=args.warmup_epochs,
            cycle_limit=getattr(args, 'lr_cycle_limit', 1),
            t_in_epochs=True,
            noise_range_t=noise_range,
            noise_pct=getattr(args, 'lr_noise_pct', 0.67),
            noise_std=getattr(args, 'lr_noise_std', 1.),
            noise_seed=getattr(args, 'seed', 42),
        )
        num_epochs = lr_scheduler.get_cycle_length() + args.cooldown_epochs
    elif args.sched == 'tanh':
        lr_scheduler = TanhLRScheduler(
            optimizer,
            t_initial=num_epochs,
            t_mul=getattr(args, 'lr_cycle_mul', 1.),
            lr_min=args.min_lr,
            warmup_lr_init=args.warmup_lr,
            warmup_t=args.warmup_epochs,
            cycle_limit=getattr(args, 'lr_cycle_limit', 1),
            t_in_epochs=True,
            noise_range_t=noise_range,
            noise_pct=getattr(args, 'lr_noise_pct', 0.67),
            noise_std=getattr(args, 'lr_noise_std', 1.),
            noise_seed=getattr(args, 'seed', 42),
        )
        num_epochs = lr_scheduler.get_cycle_length() + args.cooldown_epochs
    elif args.sched == 'step':
        lr_scheduler = StepLRScheduler(
            optimizer,
            decay_t=args.decay_epochs,
            decay_rate=args.decay_rate,
            warmup_lr_init=args.warmup_lr,
            warmup_t=args.warmup_epochs,
            noise_range_t=noise_range,
            noise_pct=getattr(args, 'lr_noise_pct', 0.67),
            noise_std=getattr(args, 'lr_noise_std', 1.),
            noise_seed=getattr(args, 'seed', 42),
        )
    elif args.sched == 'plateau':
        mode = 'min' if 'loss' in getattr(args, 'eval_metric', '') else 'max'
        lr_scheduler = PlateauLRScheduler(
            optimizer,
            decay_rate=args.decay_rate,
            patience_t=args.patience_epochs,
            lr_min=args.min_lr,
            mode=mode,
            warmup_lr_init=args.warmup_lr,
            warmup_t=args.warmup_epochs,
            cooldown_t=0,
            noise_range_t=noise_range,
            noise_pct=getattr(args, 'lr_noise_pct', 0.67),
            noise_std=getattr(args, 'lr_noise_std', 1.),
            noise_seed=getattr(args, 'seed', 42),
        )

    return lr_scheduler, num_epochs