import copy
import datetime
import gzip
import json
import os
from collections import defaultdict
from hashlib import md5
from typing import List, Tuple

import jax
import jax.numpy as jnp
import numpy as np
from omegaconf import OmegaConf

import wandb
from kinetix.environment.env_state import EnvParams, StaticEnvParams
from kinetix.environment.spaces import ActionType
from kinetix.environment.spaces import ObservationType
from kinetix.environment.ued.ued_state import UEDParams
from kinetix.util.saving import get_correct_path_of_json_level, load_from_json_file


def get_hash_without_seed(config):
    old_seed = config["seed"]
    config["seed"] = 0
    ans = md5(OmegaConf.to_yaml(config, sort_keys=True).encode()).hexdigest()
    config["seed"] = old_seed
    return ans


def get_date() -> str:
    return datetime.datetime.now().strftime("%Y-%m-%d-%H-%M-%S")


def generate_params_from_config(config):
    if config.get("env_size_type", "predefined") == "custom":
        # must load env params from a file
        _, static_env_params, env_params = load_from_json_file(config["custom_path"])
        return env_params, static_env_params.replace(
            frame_skip=config["frame_skip"],
        )
    env_params = EnvParams().replace(dense_reward_scale=config["dense_reward_scale"])

    static_env_params = StaticEnvParams().replace(
        num_polygons=config["num_polygons"],
        num_circles=config["num_circles"],
        num_joints=config["num_joints"],
        num_thrusters=config["num_thrusters"],
        frame_skip=config["frame_skip"],
        num_motor_bindings=config["num_motor_bindings"],
        num_thruster_bindings=config["num_thruster_bindings"],
    )

    return env_params, static_env_params


def get_eval_level_groups(eval_levels: List[str]) -> List[Tuple[str, str]]:
    def get_groups(s):
        # This is the size group
        group_one = s.split("/")[0]
        group_two = s.split("/")[1].split("_")[0]
        group_two = "".join([i for i in group_two if not i.isdigit()])
        if group_two == "h":
            group_two = "handmade"
        if group_two == "r":
            group_two = "random"
        return f"{group_one}_all", f"{group_one}_{group_two}"

    indices = defaultdict(list)

    for idx, s in enumerate(eval_levels):
        groups = get_groups(s)
        for group in groups:
            indices[group].append(idx)

    indices2 = {}
    for g in indices:
        indices2[g] = np.array(indices[g])

    return indices2


def normalise_config(config, name, editor_config=False, save_config: bool = True):
    old_config = copy.deepcopy(config)
    keys = ["env", "learning", "model", "misc", "eval", "ued", "env_size", "train_levels"]
    for k in keys:
        if k not in config:
            config[k] = {}
        small_d = config[k]
        del config[k]
        for kk, vv in small_d.items():
            assert kk not in config, kk
            config[kk] = vv

    config["observation_type_str"] = config["observation_type"]
    config["observation_type"] = ObservationType.from_string(config["observation_type"])
    config["action_type_str"] = config["action_type"]
    config["action_type"] = ActionType.from_string(config["action_type"])
    if not editor_config:
        config[
            "env_name"
        ] = f"Kinetix-{config['observation_type_str'].title().replace('_', '')}-{config['action_type_str'].title().replace('_', '')}"
        if config["num_train_envs"] == 2048 and config["observation_type"] == ObservationType.PIXELS:
            config["num_train_envs"] = 512
        if "SFL" in name and config["env_size_name"] in ["m", "l"]:
            config["eval_num_attempts"] = 6  # to avoid a very weird XLA bug.
        config["hash"] = get_hash_without_seed(config)

        config["random_hash"] = np.random.randint(2**31)

        config["log_save_path"] = f"logs/{config['hash']}/{config['seed']}-{get_date()}"
        if save_config:
            os.makedirs(config["log_save_path"], exist_ok=True)
            with open(f"{config['log_save_path']}/config.yaml", "w") as f:
                f.write(OmegaConf.to_yaml(old_config))
        if config["group"] == "auto":
            config["group"] = f"{name}-" + config["group_auto_prefix"] + config["env_name"].replace("Kinetix-", "")
            config["group"] += "-" + str(config["env_size_name"])

        if config["eval_levels"] == ["auto"] or config["eval_levels"] == "auto":
            config["eval_levels"] = config["train_levels_list"]
            print("Using Auto eval levels:", config["eval_levels"])
        config["num_eval_levels"] = len(config["eval_levels"])

        steps = (
            config["num_steps"]
            * config.get("outer_rollout_steps", 1)
            * config["num_train_envs"]
            * (2 if name == "PAIRED" else 1)
        )
        config["num_updates"] = int(config["total_timesteps"]) // steps

        nsteps = int(config["total_timesteps"] // 1e6)
        letter = "M"
        if nsteps >= 1000:
            nsteps = nsteps // 1000
            letter = "B"
        config["run_name"] = (
            config["env_name"] + f"-{name}-" + str(nsteps) + letter + "-" + str(config["num_train_envs"])
        )

        if config["checkpoint_save_freq"] >= config["num_updates"]:
            config["checkpoint_save_freq"] = config["num_updates"]
    return config


def get_tags(config, name):
    return [name]
    tags = [name]
    if name in ["PLR", "ACCEL", "DR"]:
        if config["use_accel"]:
            tags.append("ACCEL")
        else:
            tags.append("PLR")
    return tags


def init_wandb(config, name) -> wandb.run:
    run = wandb.init(
        config=config,
        project=config["wandb_project"],
        group=config["group"],
        name=config["run_name"],
        entity=config["wandb_entity"],
        mode=config["wandb_mode"],
        tags=get_tags(config, name),
    )
    wandb.define_metric("timing/num_updates")
    wandb.define_metric("timing/num_env_steps")
    wandb.define_metric("*", step_metric="timing/num_env_steps")
    wandb.define_metric("timing/sps", step_metric="timing/num_env_steps")
    return run


def save_data_to_local_file(data_to_save, config):
    if not config.get("save_local_data", False):
        return

    def reverse_in(li, value):
        for i, v in enumerate(li):
            if v in value:
                return True
        return False

    clean_data = {k: v for k, v in data_to_save.items() if not reverse_in(["media/", "images/"], k)}

    def _clean(x):
        if isinstance(x, jnp.ndarray):
            return x.tolist()
        elif isinstance(x, jnp.float32):
            if jnp.isnan(x):
                return -float("inf")
            return round(float(x) * 1000) / 1000
        elif isinstance(x, jnp.int32):
            return int(x)
        return x

    clean_data = jax.tree.map(lambda x: _clean(x), clean_data)
    print("Saving this data:", clean_data)
    with open(f"{config['log_save_path']}/data.jsonl", "a+") as f:
        f.write(json.dumps(clean_data) + "\n")


def compress_log_files_after_run(config):
    fpath = f"{config['log_save_path']}/data.jsonl"
    with open(fpath, "rb") as f_in, gzip.open(fpath + ".gz", "wb") as f_out:
        f_out.writelines(f_in)


def get_video_frequency(config, update_step):
    frac_through_training = update_step / config["num_updates"]
    vid_frequency = (
        config["eval_freq"]
        * config["video_frequency"]
        * jax.lax.select(
            (0.1 <= frac_through_training) & (frac_through_training < 0.3),
            1,
            jax.lax.select(
                (0.3 <= frac_through_training) & (frac_through_training < 0.6),
                2,
                4,
            ),
        )
    )
    return vid_frequency
