'''
Auxiliary functions to analyse results obtained using scripts in
the ../misc directory.

This needs to be kept in sync with the scripts and qpms/argproc.py.
'''
import math
import numpy as np
import qpms
from qpms.qpms_c import Particle
from qpms.cytmatrices import CTMatrix, TMatrixGenerator
from qpms.cybspec import BaseSpec
from qpms.cymaterials import EpsMu, EpsMuGenerator, LorentzDrudeModel, lorentz_drude
from qpms.symmetries import point_group_info
from qpms import FinitePointGroup, ScatteringSystem, eV, hbar
eh = eV / hbar

def cleanarray(a, atol=1e-8, copy=True, fullNaN=True):
    a = np.array(a, copy=copy)
    sieve = abs(a.real) < atol
    a[sieve] = 1j * a[sieve].imag
    sieve = abs(a.imag) < atol
    a[sieve] = a[sieve].real
    if fullNaN:
        if np.all(a == 0):
            a[...] = 0
    return a

def nicerot(a, atol=1e-10, copy=True): 
    '''
    Gives an array a "nice" phase.
    '''
    a = np.array(a, copy=copy)
    i = np.argmax(abs(a))
    a = a / a[i] * abs(a[i])
    return a

def _meta_matspec2emg(matspec):
    if isinstance(matspec, (float, complex)): # number is interpreted as refractive index
        return EpsMuGenerator(EpsMu(matspec**2)) 
    else:
        return EpsMuGenerator(lorentz_drude[matspec])


class FiniteRectLatAnalysis:
    '''Recreates the scattering system structure from finiterectlat-modes.py or 
    finiterectlat-constant-driving.py'''
    def __init__(self, data):
        meta = data['meta'][()]
        scatter = 'symmetry_adapted' in meta.keys()# finiterectlat-constant-driving ; TODO unify
        if scatter: 
            #thegroup = meta['symmetry_adapted'] # always D2h, this option does something else!
            thegroup = 'D2h'
        else: # finiterectlat-modes
            thegroup = 'D2h' if meta['D2'] else 'D4h'
        Nx, Ny = meta['size']
        px, py = meta['period']
        orig_x = (np.arange(Nx/2) + (0 if (Nx % 2) else .5)) * px
        orig_y = (np.arange(Ny/2) + (0 if (Ny % 2) else .5)) * py
        orig_xy = np.stack(np.meshgrid(orig_x, orig_y), axis = -1)
        bspec = BaseSpec(lMax = meta['lMax'])
        nelem = len(bspec)
        bg = _meta_matspec2emg(meta['background'])
        fg = _meta_matspec2emg(meta['material'])
        if meta.get('height', None) is None:
            tmgen = TMatrixGenerator.sphere(bg, fg, meta['radius'])
        else:
            tmgen = TMatrixGenerator.cylinder(bg, fg, meta['radius'], meta['height'], lMax_extend=meta['lMax_extend'])
        particles = [Particle(orig_xy[i], tmgen, bspec) for i in np.ndindex(orig_xy.shape[:-1])]
        sym = FinitePointGroup(point_group_info[thegroup])
        omega = data['omega'][()] if scatter else meta['centre'] * eh 
        self.ss, self.ssw = ScatteringSystem.create(particles, bg, omega, sym=sym)
        ss1, ssw1 = ScatteringSystem.create([Particle([0,0,0], tmgen, bspec)], bg, omega, sym=sym)
        self.ss1, self.ssw1 = ss1, ssw1
        
        # per-particle projectors/transformation to symmetry-adapted bases
        fvcs1 = np.empty((nelem, nelem), dtype=complex)
        y = 0
        iris1 = []
        for iri1 in range(ss1.nirreps):
            for j in range(ss1.saecv_sizes[iri1]):
                pvc1 = np.zeros((ss1.saecv_sizes[iri1],), dtype=complex)
                pvc1[j] = 1
                fvcs1[y] = ss1.unpack_vector(pvc1, iri1)
                fvcs1[y] = cleanarray(nicerot(fvcs1[y], copy=False), copy=False)
                y += 1
                iris1.append(iri1)
        
        # Mapping between ss particles and grid positions
        positions = self.ss.positions.astype(np.float32) 
        # round to get rid of duplicities due to rounding errors
        positions = positions.round(7-int(math.log(np.amax(abs(positions)),10)))
        xpositions = np.unique(positions[:,0]) 
        ypositions = np.unique(positions[:,1])
        assert(len(xpositions) == Nx)
        assert(len(ypositions == Ny))
        # particle positions as integer indices
        posmap = np.empty((positions.shape[0],2), dtype=int)
        #invposmap = np.empty((Nx, Ny), dtype=int)
        for i, pos in enumerate(positions):
            posmap[i,0] = np.searchsorted(xpositions, positions[i,0])
            posmap[i,1] = np.searchsorted(ypositions, positions[i,1])
            #invposmap[posmap[i,0], posmap[i, 1]] = i

        self.meta = meta
        self.npzfile = data
        self.fvcs1 = fvcs1
        self.iris1 = np.array(iris1)
        self.bspec = bspec
        self.nelem = nelem
        self.Nx, self.Ny, self.px, self.py = Nx, Ny, px, py
        self.posmap = posmap
        self.fullvec_poffsets = nelem * np.arange(Nx*Ny)
        self.group = thegroup
        if scatter:
            self.typ = 'scatter'
            self.fvcs1_fromdata = data['fvcs1']
            self.iris1_fromdata = data['iris1']
            self.positions_fromdata = data['positions']
            self.driving_symmetry_adapted = data['symmetry_adapted']
        else: # modes
            self.typ = 'modes'
            self.eigval = data['eigval']
            self.neig = len(self.eigval)
            self.eigvec = data['eigvec']
            self.residuals = data['residuals']
            self.ranktest_SV = data['ranktest_SV']
            self.iri = data['iri'][()]


    def fullvec2grid(self, fullvec, swapxy=False):
        arr = np.empty((self.Nx, self.Ny, self.nelem), dtype=complex)
        for pi, offset in enumerate(self.fullvec_poffsets):
            ix, iy = self.posmap[pi]
            arr[ix, iy] = fullvec[offset:offset+self.nelem]
        return np.swapaxes(arr, 0, 1) if swapxy else arr
    
    def saecv2grid(self, saecv, swapxy=False):
        fullvec = self.ss.unpack_vector(saecv, iri=self.iri)
        return self.fullvec2grid(fullvec, swapxy)
    
    def eigvec_asgrid(self, i, swapxy = False):
        '''
        Returns i-th eigenmode as a (Nx, Ny, nelem)-shaped grid of excitation
        coefficients.
        '''
        if self.iri is not None:
            return self.saecv2grid(self.eigvec[i], swapxy)
        else:
            return self.fullvec2grid(self.eigvec[i], swapxy)