# !==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==
# >-|===|>                             Imports                             <|===|-<
# !==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==
from kgsim.environment import dHybridRtemplate, frameDir
from kgsim.fields import ScalarField, VectorField
from kgsim.simulation import GenericSimulation, SimulationGroup
from kgsim.particles import Species, Particle
from kgsim.dhybridr.io import dHybridRinput, dHybridRout
from kgsim.dhybridr.initializer import dHybridRinitializer, TurbInit, dHybridRconfig
from kgsim.dhybridr.anvil_submit import AnvilSubmitScript

from kplot import show, func_video
from kbasic.Tex import texfraction
from kbasic.parsing import Folder
from kbasic.user_input import yesno
from os import system
from os.path import isdir
from glob import glob

from matplotlib.pyplot import cm as cmaps
import numpy as np 
from h5py import File as h5File

# !==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==
# >-|===|>                           Definitions                           <|===|-<
# !==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==
track_keys = ['B1', 'B2', 'B3', 'E1', 'E2', 'E3', 'ene', 'n', 'p1', 'p2', 'p3', 'q', 't', 'x1', 'x2', 'x3']

# !==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==
# >-|===|>                            Functions                            <|===|-<
# !==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==
# simulation parsing
def extract_energy(file_name: str) -> tuple:
    with h5File(file_name, 'r') as file:
        fE = np.mean(file["DATA"], axis=1)
        [low, high] = file["AXIS"]["X2 AXIS"][:]
        lne = np.linspace(low, high, fE.shape[0])
        dlne = np.diff(lne)[0]
        E = np.exp(lne)
        return E, fE, dlne
def iters(simulation) -> list[int]:
    return [int(fn[-11:-3]) for fn in simulation.density.file_names]
def times(simulation) -> list[float]:
    return list(np.array(iters(simulation)).astype(float) * simulation.dt)
def particle_video(
    sim,
    particles: list[str] | int,
    fname: str,
    background = None,      
    resume: bool = False,
    res: int = None,
    zfill: int = 10,
    dpi: int = 250,
    # paticle plotting keywords
    color = 'red',
    marker='.',
    ms=10, 
    # assume the rest are keywords for the show function
    **kwds
):
    # check species
    assert sim.input.num_species == 1, "only one species is implemented rn :-("
    # check the time resolution is valid
    if not res: res = sim.input.sp01.track_nstore
    dnp = sim.input.sp01.track_nstore
    dnb = sim.input.ndump
    assert (dnp % res == 0) & (res % dnb == 0), "your resolution must be an integer multiple of track_nstore and ndump must be an integer multiple of res."
    # handle the resuming
    if not resume: 
        system(f'rm {frameDir.path}/*')
        i_start = 0
    else:
        last_file = sorted(frameDir.children)[-1]
        last_iter = int(last_file[:-4])
        i_start = last_iter + res
    # find the file
    fn = sim.path+"/Output/Tracks/Sp01/track_Sp01.h5"
    with h5File(fn) as file:
        # extract particle tracks based on the particles argument
        tags = np.array(list(file.keys()))
        match particles:
            case [str(x), ]: selected = particles 
            case int(x): selected = np.random.choice(tags, size=x, replace=False)
        sx, sy = np.array([file[t]['x1'] for t in selected]).T, np.array([file[t]['x2'] for t in selected]).T 
        # setup background
        if not background: background = sim.density
        # initialize plots
        fig,ax,img = show(
            background[0], 
            x=range(0, sim.input.boxsize[0], sim.dx), y=range(0, sim.input.boxsize[1], sim.dy), 
            zorder=1,
            **kwds
        )
        line, = ax.plot(
            sx[0], sy[0], 
            ls='None', color=color, marker=marker, ms=ms,
            zorder = 2
        )
        # execute the loop
        def update(i: int, dnp=dnp, dnb=dnb):
            pind = i * dnp
            line.set_data(sx[pind], sy[pind])
            if pind % dnb == 0:
                bind = pind // dnb 
                img.set_array(background[bind])
        func_video(fname, fig, update, )

# !==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==
# >-|===|>                             Classes                             <|===|-<
# !==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==
class dHybridRparticle(Particle):
    def __init__(self, tag: str, species: Species, load: bool = False):
        self.tag = tag
        self.species = self.sp = species
        if load: self.load()
    def load(self):
        with h5File(self.species.path) as file:
            for k in track_keys:
                try: setattr(self, k, file[self.tag][k][:])
                except: continue
            self.x = self.x1 
            self.y = self.x2
    def follow_video(self, background=None, window=100, **kwds):
        sim = self.species.parent
        if not background: background = sim.density
        di = sim.input.sp01.track_nstore
        db = sim.input.ndump
        center = (window//2, window//2)
        fig,ax,img = show(background[0], **kwds)
        line, = ax.plot(*center, ls='None', marker='.', ms=10, color='r')
        for i in range(0, sim.input.niter, di):
            pind = i // di 
            img.set
class dHybridRspecies(Species):
    def __init__(self, species_number: int, parent):
        self.n = self.number = self.sp_num = species_number
        self.nstr = str(self.n).zfill(2)
        self.parent = parent 
        self.input = getattr(parent.input, f"sp{str(self.n).zfill(2)}")
        self.loaded = False
        if self.input.track_dump:
            self.path = parent.path+f"/Output/Tracks/Sp{self.nstr}/track_Sp{self.nstr}.h5"
            self.load()
    
    def load(self):
        with h5File(self.path) as file:
            self.tags = np.array(list(file.keys()))
            self.loaded = True 

    def sample(self, N, load=False): return np.array([dHybridRparticle(t, self) for t in np.random.choice(self.tags, size=N, load=load)])
class dHybridR(GenericSimulation):
    """
    A simulation class to interact with dHybridR simulations in python
    """
    def __init__(
            self, 
            path: str,
            caching: bool = False,
            verbose: bool = False,
            template: Folder = dHybridRtemplate,
            compressed: bool = False
        ) -> None:
        self.compressed = compressed
        #setup simulation
        GenericSimulation.__init__(self, path, caching=caching, verbose=verbose, template=template)
        #setup input, output, and restart folders
        self.parse_input()
        self.output = Folder(self.path+"/Output")
        if not self.output.exists and verbose:
            if yesno("There is no output, would you like to run this simulation?\n"): 
                self.run()
        elif self.output.exists: 
            self.parse_output()
            self.ncores: int = int(np.prod(self.input.node_number))
            self.ncores_charged: int = self.ncores + self.ncores % 128
            if self.out.exists:
                self.runtime: float = self.out.runtime #run time as calculated from out file, in hours
                self.corehours: float = self.runtime * np.prod(self.ncores)
                self.corehours_charged: float = self.runtime * self.ncores_charged
        self.restartDir = Folder(self.path+"/Restart")
    def __repr__(self) -> str: return self.name
    def __len__(self) -> int:
        if self.output.exists: 
            return len(self.density)
        else: return 0
    def create(self) -> None:
        self.template.copy(self.path)
        system(f"chmod 755 {self.path}/dHybridR")
    def parse_input(self) -> None:
        self.input = dHybridRinput(self.path+"/input/input")
        self.dt = self.input.dt
        self.niter = self.input.niter
        self.dx: float = self.input.boxsize[0]/self.input.ncells[0]
        self.dy: float = self.input.boxsize[1]/self.input.ncells[1]
        self.dz: float = self.input.boxsize[2]/self.input.ncells[2] if len(self.input.boxsize)==3 else np.inf
    def run(self, initializer: dHybridRinitializer, submit_script: AnvilSubmitScript) -> None:
        initializer.prepare_simulation()
        submit_script.write()
        system(f"sh {submit_script.path}")
    def parse_output(self) -> None:
        self.out = dHybridRout(self.path+"/out")
        kwargs = {'caching':self.caching, 'verbose':self.verbose, 'parent':self}
        self.B       = VectorField(self.path + "/Output/Fields/Magnetic/Total/", name="magnetic", latex="B", **kwargs)
        self.E       = VectorField(self.path + "/Output/Fields/Electric/Total/", name="electric", latex="E", **kwargs)
        self.etx1    = ScalarField(self.path + "/Output/Phase/etx1/Sp01/", **kwargs)
        self.pxx1    = ScalarField(self.path + "/Output/Phase/p1x1/Sp01/", **kwargs)
        self.pyx1    = ScalarField(self.path + "/Output/Phase/p2x1/Sp01/", **kwargs)
        self.pzx1    = ScalarField(self.path + "/Output/Phase/p3x1/Sp01/", **kwargs)
        self.density = ScalarField(self.path + "/Output/Phase/x3x2x1/Sp01/", name="density", latex=r"$\rho$", **kwargs)
        self.Pxx     = ScalarField(self.path + "/Output/Phase/PressureTen/Sp01/xx/", **kwargs)
        self.Pyy     = ScalarField(self.path + "/Output/Phase/PressureTen/Sp01/yy/", **kwargs)
        self.Pzz     = ScalarField(self.path + "/Output/Phase/PressureTen/Sp01/zz/", **kwargs)
        self.u       = VectorField(self.path + "/Output/Phase/FluidVel/Sp01/", name="bulkflow", latex="u", **kwargs)
        [
            self.energy_grid,
            self.energy_pdf,
            self.dlne
        ] = np.array([[*extract_energy(f)] for f in self.etx1.file_names], dtype=object).T
        self.energy_grid = np.vstack(self.energy_grid) 
        self.energy_pdf = np.vstack(self.energy_pdf) 
        self.dlne = np.vstack(self.dlne) 
        if self.input.sp01.track_dump:
            self.sp01 = dHybridRspecies(1, self)
        self.iter = iters(self)
        self.time = times(self)
class dHybridRgroup(SimulationGroup):
    def __init__(self, path, **sim_kwds):
        SimulationGroup.__init__(self, path, simtype=dHybridR, **sim_kwds)