import os
import uuid
from dataclasses import asdict, dataclass
from typing import Any, DefaultDict, Dict, List, Optional, Tuple

import bullet_safety_gym  # noqa
import dsrl
import gymnasium as gym  # noqa
import numpy as np
import pyrallis
import torch
from dsrl.infos import DENSITY_CFG
from dsrl.offline_env import OfflineEnvWrapper, wrap_env  # noqa
from fsrl.utils import WandbLogger
from torch.utils.data import DataLoader
from tqdm.auto import trange  # noqa
import sys
sys.path.append("../A-Sepsis-RL/CDT")

from examples_my_cost.configs.coptidice_configs import (COptiDICE_DEFAULT_CONFIG,
                                                COptiDICETrainConfig)
from osrl_my_cost.algorithms import COptiDICE, COptiDICETrainer
from osrl_my_cost.common import TransitionDataset
from osrl_my_cost.common.exp_util import auto_name, seed_all
import pickle

@pyrallis.wrap()
def train(args: COptiDICETrainConfig):
    seed_all(args.seed)
    if args.device == "cpu":
        torch.set_num_threads(args.threads)

    # setup logger
    cfg = asdict(args)
    default_cfg = asdict(COptiDICE_DEFAULT_CONFIG[args.task]())
    if args.name is None:
        args.name = auto_name(default_cfg, cfg, args.prefix, args.suffix)
    if args.group is None:
        args.group = args.task + "-cost-" + str(int(args.cost_limit))
    if args.logdir is not None:
        args.logdir = os.path.join(args.logdir, args.group, args.name)
    logger = WandbLogger(cfg, args.project, args.group, args.name, args.logdir)
    # logger = TensorboardLogger(args.logdir, log_txt=True, name=args.name)
    logger.save_config(cfg, verbose=args.verbose)

    # initialize environment
    # env = gym.make(args.task)

    # pre-process offline dataset
    # data = env.get_dataset()
    # env.set_target_cost(args.cost_limit)

    # dataset_path_val = f'./CDT/examples_my_cost/data/my_cdt_data_val_noauto_c6.pkl'
    # dataset_path = f'./CDT/examples_my_cost/data/my_cdt_data_train_noauto_c6.pkl'

    # dataset_path_val = f'/home/fn/A-Sepsis-RL/CDT/examples_my_cost/data/transicrl_data_val.pkl'
    # dataset_path = f'/home/fn/A-Sepsis-RL/CDT/examples_my_cost/data/transicrl_data.pkl'

    dataset_path_val = f'./CDT/examples_my_cost/data/cost0_data_val.pkl'
    dataset_path = f'./CDT/examples_my_cost/data/cost0_data_val.pkl'

    with open(dataset_path, 'rb') as f:
        data = pickle.load(f)

    with open(dataset_path_val,'rb') as f:
        data_val = pickle.load(f)

    cbins, rbins, max_npb, min_npb = None, None, None, None
    if args.density != 1.0:
        density_cfg = DENSITY_CFG[args.task + "_density" + str(args.density)]
        cbins = density_cfg["cbins"]
        rbins = density_cfg["rbins"]
        max_npb = density_cfg["max_npb"]
        min_npb = density_cfg["min_npb"]
    # data = env.pre_process_data(data,
    #                             args.outliers_percent,
    #                             args.noise_scale,
    #                             args.inpaint_ranges,
    #                             args.epsilon,
    #                             args.density,
    #                             cbins=cbins,
    #                             rbins=rbins,
    #                             max_npb=max_npb,
    #                             min_npb=min_npb)

    # wrapper
    # env = wrap_env(
    #     env=env,
    #     reward_scale=args.reward_scale,
    # )
    # env = OfflineEnvWrapper(env)

    # setup dataset

    state_dim = 48
    action_dim = 2
    max_action = [1.0,1.0]

    dataset = TransitionDataset(data,
                                reward_scale=args.reward_scale,
                                cost_scale=args.cost_scale,
                                state_init=True)
    trainloader = DataLoader(
        dataset,
        batch_size=args.batch_size,
        pin_memory=True,
        num_workers=args.num_workers,
    )
    trainloader_iter = iter(trainloader)
    init_s_propotion, obs_std, act_std = dataset.get_dataset_states()

    # setup model
    model = COptiDICE(
        state_dim=state_dim,
        action_dim=action_dim,
        max_action=max_action,
        f_type=args.f_type,
        init_state_propotion=init_s_propotion,
        observations_std=obs_std,
        actions_std=act_std,
        a_hidden_sizes=args.a_hidden_sizes,
        c_hidden_sizes=args.c_hidden_sizes,
        gamma=args.gamma,
        alpha=args.alpha,
        cost_ub_epsilon=args.cost_ub_epsilon,
        num_nu=args.num_nu,
        num_chi=args.num_chi,
        cost_limit=args.cost_limit,
        episode_len=args.episode_len,
        device=args.device,
    )

    print(f"Total parameters: {sum(p.numel() for p in model.parameters())}")

    def checkpoint_fn():
        return {"model_state": model.state_dict()}

    logger.setup_checkpoint_fn(checkpoint_fn)

    trainer = COptiDICETrainer(model,
                               logger=logger,
                               actor_lr=args.actor_lr,
                               critic_lr=args.critic_lr,
                               scalar_lr=args.scalar_lr,
                               reward_scale=args.reward_scale,
                               cost_scale=args.cost_scale,
                               device=args.device)

    # for saving the best
    best_reward = -np.inf
    best_cost = np.inf
    best_idx = 0

    for step in trange(args.update_steps, desc="Training"):
        batch = next(trainloader_iter)
        observations, next_observations, actions, rewards, costs, done, is_init = [
            b.to(args.device) for b in batch
        ]
        trainer.train_one_step(observations.float(), next_observations.float(), actions.float(), rewards.float(), costs.float(), done, is_init)

        # # evaluation
        # if (step + 1) % args.eval_every == 0 or step == args.update_steps - 1:
        #     cdt_actions, phy_actions, length = trainer.evaluate(args.eval_episodes)
        #     logger.store(tab="eval", Cost=cost, Reward=ret, Length=length)

        #     # save the current weight
        #     logger.save_checkpoint()
        #     # save the best weight
        #     if cost < best_cost or (cost == best_cost and ret > best_reward):
        #         best_cost = cost
        #         best_reward = ret
        #         best_idx = step
        #         logger.save_checkpoint(suffix="best")

        #     logger.store(tab="train", best_idx=best_idx)
        #     logger.write(step, display=False)
        logger.write_without_reset(step)

    path_ =f'/home/fn/A-Sepsis-RL/CDT/Mymodel/coptidice_cost0.pt'
    torch.save(model.state_dict(),path_)

if __name__ == "__main__":
    train()
