
"""
Distributed Computing
"""


from typing import Optional, Any, Callable, Union
from math import ceil, floor
from enum import Enum
from dataclasses import dataclass
import psutil
import numpy as np
import crocodile.toolbox as tb
from crocodile.cluster.remote_machine import RemoteMachine, RemoteMachineConfig, WorkloadParams, LAUNCH_METHOD
from rich.console import Console
# from platform import system
# import time
# from rich.progress import track


console = Console()


class LoadCriterion(Enum):
    cpu = "cpu"
    ram = "ram"
    product = "cpu * ram"
    cpu_norm = "cpu_norm"
    ram_norm = "ram_norm"
    product_norm = "cpu_norm * ram_norm"


@dataclass
class MachineSpecs:
    cpu: float
    ram: float
    product: float
    cpu_norm: float
    ram_norm: float
    product_norm: float
    @staticmethod
    def get_this_machine_specs():
        cpu, ram = psutil.cpu_count(), psutil.virtual_memory().total / 2 ** 30
        return MachineSpecs(cpu=cpu, ram=ram, product=cpu * ram, cpu_norm=cpu, ram_norm=ram, product_norm=cpu * ram)


class ThreadLoadCalculator:
    """relies on relative values to a referenc machine specs.
    Runs multiple instances of code per machine. Useful if code doesn't run faster with more resources avaliable.
    equal distribution across instances of one machine"""
    def __init__(self, num_jobs: Optional[int] = None, load_criterion: LoadCriterion = LoadCriterion.cpu, reference_specs: Optional[MachineSpecs] = None):
        self.num_jobs = num_jobs
        self.load_criterion = load_criterion
        self.reference_specs: MachineSpecs = MachineSpecs.get_this_machine_specs() if reference_specs is None else reference_specs
    def __getstate__(self): return self.__dict__
    def __setstate__(self, state: dict[str, Any]): self.__dict__.update(state)
    def get_num_threads(self, machine_specs: MachineSpecs) -> int:
        if self.num_jobs is None: return 1
        res = int(floor(self.num_jobs * (machine_specs.__dict__[self.load_criterion.name] / self.reference_specs.__dict__[self.load_criterion.name])))
        return 1 if res == 0 else res


class MachineLoadCalculator:
    def __init__(self, max_num: int = 1000, load_criterion: LoadCriterion = LoadCriterion.product, load_ratios_repr: str = ""):
        self.load_ratios: list[float] = []
        self.load_ratios_repr = load_ratios_repr
        self.max_num: int = max_num
        self.load_criterion = load_criterion
    def __getstate__(self) -> dict[str, Any]: return self.__dict__
    def __setstate__(self, d: dict[str, Any]) -> None: self.__dict__.update(d)
    def get_workload_params(self, machines_specs: list[MachineSpecs], threads_per_machine: list[int]) -> list[WorkloadParams]:
        """Note: like thread divider in parallelize function, the behaviour is to include the edge cases on both ends of subsequent intervals."""
        tmp: list[WorkloadParams] = []
        idx_so_far = 0
        for machine_index, (machine_specs, a_threads_per_machine) in enumerate(zip(machines_specs, threads_per_machine)):
            load_value = machine_specs.__dict__[self.load_criterion.name]
            self.load_ratios.append(load_value)
            idx1 = idx_so_far
            idx2 = self.max_num if machine_index == len(threads_per_machine) - 1 else (floor(load_value * self.max_num) + idx1)
            if idx2 > self.max_num:
                print(machines_specs, '\n\n', threads_per_machine)
                print(f"All values: {tmp=}, {idx_so_far=}, {idx1=}, {idx2=}, {self.max_num=}, {a_threads_per_machine=}, {machine_index=}, {machine_specs=}, {load_value=}, {self.load_ratios=}, {self.load_ratios_repr=}")
                raise ValueError(f"idx2 ({idx2}) > max_num ({self.max_num})")
            idx_so_far = idx2
            tmp.append(WorkloadParams(idx_start=idx1, idx_end=idx2, idx_max=self.max_num, jobs=a_threads_per_machine))
        return tmp


class Cluster:
    def __getstate__(self) -> dict[str, Any]:
        state = self.__dict__
        state["func"] = None
        return state
    def __setstate__(self, state: dict[str, Any]) -> None: self.__dict__.update(state)
    def save(self) -> tb.P:
        path = self.root_dir.joinpath("cluster.Cluster.pkl")
        tb.Save.vanilla_pickle(obj=self.__getstate__(), path=path)
        return path
    @staticmethod
    def load(job_id: str, base: Optional[str] = None) -> 'Cluster': return Cluster.get_cluster_path(job_id=job_id, base=base).joinpath("cluster.Cluster.pkl").readit()
    @staticmethod
    def get_cluster_path(job_id: str, base: Union[str, tb.P, None] = None):
        if base is None: base_obj = tb.P.home().joinpath(rf"tmp_results/remote_machines")
        else: base_obj = tb.P(base)
        return base_obj.joinpath(f"job_id__{job_id}")
    def __init__(self,
                 func: Callable[..., Any],
                 ssh_params: list[dict[str, Any]],
                 remote_machine_config: RemoteMachineConfig,
                 func_kwargs: Optional[dict[str, Any]] = None,
                 # workload_params: list[WorkloadParams] or None = None,
                 thread_load_calc: Optional[ThreadLoadCalculator] = None,
                 # machine_load_calc=None,
                 ditch_unavailable_machines: bool = False,
                 description: str = "",
                 job_id: Optional[str] = None,
                 base_dir: Union[str, tb.P, None] = None):
        self.job_id = job_id or tb.randstr(noun=True)
        self.root_dir = self.get_cluster_path(self.job_id, base=base_dir)
        self.results_downloaded = False

        self.thread_load_calc: ThreadLoadCalculator = thread_load_calc or ThreadLoadCalculator()
        self.machine_load_calc: MachineLoadCalculator = MachineLoadCalculator(load_criterion=LoadCriterion[self.thread_load_calc.load_criterion.name + "_norm"], )

        sshz: list[tb.SSH] = []
        for an_ssh_params in ssh_params:
            try:
                tmp = tb.SSH(**an_ssh_params)
                sshz.append(tmp)
            except Exception as ex:
                print(f"Couldn't connect to {an_ssh_params}")
                if ditch_unavailable_machines: continue
                else: raise Exception(f"Couldn't connect to {an_ssh_params}") from ex  # type: ignore # pylint: disable=W0719

        # lists of similar length:
        self.sshz: list[tb.SSH] = sshz
        self.machines: list[RemoteMachine] = []
        self.machines_specs: list[MachineSpecs] = []
        self.threads_per_machine: list[int] = []
        self.remote_machine_kwargs: RemoteMachineConfig = remote_machine_config
        self.workload_params: list[WorkloadParams] = []

        self.description: str = description
        self.func = func
        self.func_kwargs = func_kwargs if func_kwargs is not None else {}

        # fire options
        self.machines_per_tab: int = 1
        self.window_number: int = 2

    def __repr__(self): return f"Cluster with following machines:\n" + "\n".join([repr(item) for item in (self.machines if self.machines else self.sshz)])
    def print_func_kwargs(self):
        print("\n" * 2)
        console.rule(title=f"kwargs of functions to be run on machines")
        for an_ssh, a_kwarg in zip(self.sshz, self.workload_params):
            tb.S(a_kwarg.__dict__).print(as_config=True, title=an_ssh.get_remote_repr())
    def print_commands(self, launch_method: LAUNCH_METHOD):
        print("\n" * 2)
        console.rule(title="Commands to run on each machine:")
        for machine in self.machines:
            print(f"{repr(machine)} ==> {machine.file_manager.get_fire_command(launch_method=launch_method)}")

    def generate_standard_kwargs(self) -> None:
        if self.workload_params:
            self.print_func_kwargs()
            print(self.workload_params, len(self.workload_params), type(self.workload_params))
            print("workload_params is not None, so not generating standard kwargs")
            return None
        cpus: list[float] = []
        rams: list[float] = []
        for an_ssh in self.sshz:
            res = an_ssh.run_py("import psutil; print(psutil.cpu_count(), psutil.virtual_memory().total)", verbose=False).op
            try: cpus.append(int(res.split(' ')[0]))
            except ValueError as ve:
                print(f"Couldn't get cpu count from {an_ssh}")
                raise ValueError(f"Couldn't get cpu count from {an_ssh.get_remote_repr()}") from ve
            rams.append(ceil(int(res.split(' ')[1]) / 2 ** 30))
        total_cpu = np.array(cpus).sum()
        total_ram = np.array(rams).sum()
        total_product = (np.array(cpus) * np.array(rams)).sum()

        self.machines_specs = [MachineSpecs(cpu=a_cpu, ram=a_ram, product=a_cpu * a_ram, cpu_norm=a_cpu / total_cpu, ram_norm=a_ram / total_ram, product_norm=a_cpu * a_ram / total_product) for a_cpu, a_ram in zip(cpus, rams)]
        self.threads_per_machine = [self.thread_load_calc.get_num_threads(machine_specs=machine_specs) for machine_specs in self.machines_specs]
        self.workload_params = self.machine_load_calc.get_workload_params(machines_specs=self.machines_specs, threads_per_machine=self.threads_per_machine)
        self.print_func_kwargs()

    def viz_load_ratios(self) -> None:
        if not self.workload_params: raise RuntimeError("func_kwargs_list is None. You need to run generate_standard_kwargs() first.")
        plt = tb.install_n_import("plotext")
        names = tb.L(self.sshz).apply(lambda x: x.get_remote_repr(add_machine=True)).list

        plt.simple_multiple_bar(names, [[machine_specs.cpu for machine_specs in self.machines_specs], [machine_specs.ram for machine_specs in self.machines_specs]], title=f"Resources per machine", labels=["#cpu threads", "memory size"])
        plt.show()
        print("")
        plt.simple_bar(names, self.machine_load_calc.load_ratios, width=100, title=f"Load distribution for machines using criterion `{self.machine_load_calc.load_criterion}`")
        plt.show()

        tmp = tb.S(dict(zip(names, tb.L((np.array(self.machine_load_calc.load_ratios) * 100).round(1)).apply(lambda x: f"{int(x)}%")))).print(as_config=True, justify=75, return_str=True)
        assert isinstance(tmp, str)
        self.machine_load_calc.load_ratios_repr = tmp
        print(self.machine_load_calc.load_ratios_repr)
        # self.workload_params.
        print("\n")

    def submit(self) -> None:
        if not self.workload_params: raise RuntimeError("You need to generate standard kwargs first.")
        for idx, (a_workload_params, an_ssh) in enumerate(zip(self.workload_params, self.sshz)):
            desc = self.description + f"\nLoad Ratios on machines:\n{self.machine_load_calc.load_ratios_repr}"
            # if self.remote_machine_kwargs is not None:
            config = self.remote_machine_kwargs
            config.__dict__.update(dict(description=desc, job_id=self.job_id + f"_{idx}", base_dir=self.root_dir, workload_params=a_workload_params, ssh_obj=an_ssh))
            # else: config = RemoteMachineConfig(description=desc, job_id=self.job_id + f"_{idx}", base_dir=self.root_dir.as_posix(), workload_params=a_workload_params, ssh_obj=an_ssh)
            m = RemoteMachine(func=self.func, func_kwargs=self.func_kwargs, config=config)
            m.generate_scripts()
            m.submit()
            self.machines.append(m)
        try: self.save()
        except Exception as re:
            print(re)
            print("Couldn't pickle cluster object")
        # self.print_commands()

    def open_mux(self, machines_per_tab: int = 1, window_number: Optional[int] = None):
        self.machines_per_tab = machines_per_tab
        self.window_number = window_number if window_number is not None else 0  # tb.randstr(length=3, lower=False, upper=False)
        cmd = f"wt -w {self.window_number} "
        for idx, m in enumerate(self.machines):

            sub_cmd = m.get_session_manager().get_new_session_ssh_command(ssh=m.ssh, sess_name=m.job_params.session_name)
            if idx == 0: cmd += f""" new-tab --title '{str(m.ssh.hostname) + str(idx)}' pwsh -Command "{sub_cmd}" `;"""  # avoid new tabs despite being even index
            elif idx % self.machines_per_tab == 0: cmd += f""" new-tab --title {str(m.ssh.hostname) + str(idx)} pwsh -Command "{sub_cmd}" `;"""
            else: cmd += f""" split-pane --horizontal --size {1 / self.machines_per_tab} pwsh -Command "{sub_cmd}" `;"""

        print("Terminal launch command:\n", cmd)
        if cmd.endswith("`;"): cmd = cmd[:-2]
        tb.Terminal().run_async(*cmd.replace("`;", ";").split(" "))  # `; only for powershell, cmd is okay for ; as it is not a special character
        rm_last = self.machines[-1]
        rm_last.get_session_manager().asssert_session_started(ssh=rm_last.ssh, sess_name=rm_last.job_params.session_name)

    def fire(self, machines_per_tab: int = 1, window_number: Optional[int] = None, run: bool = False):
        self.open_mux(machines_per_tab=machines_per_tab, window_number=window_number)
        for m in self.machines:
            m.fire(run=run, open_console=False)

    def run(self, run: bool = False, machines_per_tab: int = 1, window_number: Optional[int] = None):
        self.generate_standard_kwargs()
        self.viz_load_ratios()
        print(self)
        self.submit()
        self.fire(run=run, machines_per_tab=machines_per_tab, window_number=window_number)
        self.save()
        return self

    def check_job_status(self) -> None: tb.L(self.machines).apply(lambda machine: machine.check_job_status())
    def download_results(self):
        if self.results_downloaded:
            print(f"All results downloaded to {self.root_dir} 🤗")
            return True
        for idx, a_m in enumerate(self.machines):
            _ = idx
            if a_m.results_path is None:
                print(f"Results are not ready for machine {a_m}.")
                print(f"Try to run `.check_job_status()` to check if the job is done and obtain results path.")
                continue
            # results_folder = tb.P(a_m.results_path).expanduser()
            if a_m.results_downloaded is False:
                print("\n")
                console.rule(f"Downloading results from {a_m}")
                print("\n")
                a_m.download_results(target=None)  # TODO another way of resolve multiple machines issue is to create a directory at downlaod_results time.
        if tb.L(self.machines).results_downloaded.to_numpy().sum() == len(self.machines):
            print(f"All results downloaded to {self.root_dir} 🤗")
            self.results_downloaded = True


if __name__ == '__main__':
    pass
