import gym
import numpy as np
import torch
import wandb
import torch.nn as nn
import time
import collections
import argparse
import pickle
import random
import d4rl
import swanlab
## method one
# In your code file,add following:
from data import create_dataloader
from decision_transformer.evaluation.evaluate_episodes import evaluate_episode, evaluate_episode_rtg
from decision_transformer.models.cross_attetion_decision_transformer import Cross_attetion_DecisionTransformer
from decision_transformer.models.mlp_bc import MLPBCModel
from decision_transformer.training.act_trainer import ActTrainer
from decision_transformer.training.trainer_cross_attetion import Cross_attetion_Trainer
from decision_transformer.training.seq_trainer import SequenceTrainer
from decision_transformer.models.bert_decision_transformer import Bert_DecisionTransformer
import minari


def discount_cumsum(x, gamma):
    discount_cumsum = np.zeros_like(x)
    discount_cumsum[-1] = x[-1]
    for t in reversed(range(x.shape[0]-1)):
        discount_cumsum[t] = x[t] + gamma * discount_cumsum[t+1]
    return discount_cumsum
def experiment(
        exp_prefix,
        variant,
):
    device = variant.get('device', 'cuda')
    log_to_swanlab = variant.get('log_to_swanlab', False)
    env_name, dataset = variant['env'], variant['dataset']
    model_type = variant['model_type']
    group_name = f'{exp_prefix}-{env_name}-{dataset}'
    exp_prefix = f'{group_name}-{random.randint(int(1e5), int(1e6) - 1)}'
    num_subgoal = variant['num_subgoal']
    temperature = variant['temperature']
    if env_name == 'hopper':
        env = gym.make(f'{env_name}-{dataset}-v0')
        max_ep_len = 1000
        env_targets = [5000,3600, 1800]  # evaluation conditioning targets
        scale = 1000.  # normalization for rewards/returns
    elif env_name == 'halfcheetah':
        env = gym.make(f'{env_name}-{dataset}-v0')
        max_ep_len = 1000
        env_targets = [12000,9000, 6000]
        scale = 1000.
    elif env_name == 'walker2d':
        env = gym.make(f'{env_name}-{dataset}-v0')
        max_ep_len = 1000
        env_targets = [8000,6000,5000,4000, 2500]
        scale = 1000.
    elif env_name == 'reacher2d':
        from decision_transformer.envs.reacher_2d import Reacher2dEnv
        env = Reacher2dEnv()
        max_ep_len = 100
        env_targets = [76, 40]
        scale = 10.
    elif env_name =='maze2d':
        env = gym.make(f'{env_name}-{dataset}-v1')
        max_ep_len = 1000
        env_targets = [500,300,200,150,100,50,30]
        scale = 100
    elif env_name == 'kitchen':
        dversion = 0
        gym_name = f'{env_name}-{dataset}-v{dversion}'
        env = gym.make(gym_name)
        max_ep_len = 1000
        env_targets = [500, 250]
        scale = 100.
    if model_type == 'bc':
        env_targets = env_targets[:1]  # since BC ignores target, no need for different evaluations

    state_dim = env.observation_space.shape[0]
    act_dim = env.action_space.shape[0]
    future_dim = variant.get('future_dim', '16')
    # load dataset
    dataset_path = f'data/data-maze2d/{env_name}-{dataset}-v1.h{num_subgoal}py'
    with open(dataset_path, 'rb') as f:
        trajectories = pickle.load(f)

    # save all path information into separate lists
    mode = variant.get('mode', 'normal')
    states, traj_lens, returns = [], [], []
    for path in trajectories:
        if mode == 'delayed':  # delayed: all rewards moved to end of trajectory
            path['rewards'][-1] = path['rewards'].sum()
            path['rewards'][:-1] = 0.
        states.append(path['observations'])
        traj_lens.append(len(path['observations']))
        returns.append(path['rewards'].sum())
    traj_lens, returns = np.array(traj_lens), np.array(returns)

    # used for input normalization
    states = np.concatenate(states, axis=0)
    state_mean, state_std = np.mean(states, axis=0), np.std(states, axis=0) + 1e-6

    num_timesteps = sum(traj_lens)

    print('=' * 50)
    print(f'Starting new experiment: {env_name} {dataset}')
    print(f'{len(traj_lens)} trajectories, {num_timesteps} timesteps found')
    print(f'Average return: {np.mean(returns):.2f}, std: {np.std(returns):.2f}')
    print(f'Max return: {np.max(returns):.2f}, min: {np.min(returns):.2f}')
    print('=' * 50)

    K = variant['K']
    batch_size = variant['batch_size']
    num_eval_episodes = variant['num_eval_episodes']
    pct_traj = variant.get('pct_traj', 1.)

    # only train on top pct_traj trajectories (for %BC experiment)
    num_timesteps = max(int(pct_traj*num_timesteps), 1)
    sorted_inds = np.argsort(returns)  # lowest to highest
    num_trajectories = 1
    timesteps = traj_lens[sorted_inds[-1]]
    ind = len(trajectories) - 2
    while ind >= 0 and timesteps + traj_lens[sorted_inds[ind]] <= num_timesteps:
        timesteps += traj_lens[sorted_inds[ind]]
        num_trajectories += 1
        ind -= 1
    sorted_inds = sorted_inds[-num_trajectories:]

    # used to reweight sampling so we sample according to timesteps instead of trajectories
    p_sample = traj_lens[sorted_inds] / sum(traj_lens[sorted_inds])
    def generate_arithmetic_sequence(k):#等差[0.8,0.6,0.4,0.2,0.0] k个sub均匀分布在[0,1]
        sequence = []
        threshold=round( (1/k) , 2)# 保留两位小数
        current_value = 1-threshold
        for i in range(k):
            sequence.append(round(current_value, 2))  
            current_value -= threshold
        return sequence
    def get_batch_new(batch_size=256,num_subgoal=5,max_len=K):
        batch_inds = np.random.choice(
            np.arange(num_trajectories),
            size=batch_size,
            replace=True,
            p=p_sample,  # reweights so we sample according to timesteps
        )
        s, target_rtg ,subgoal= [], [] ,[]
        for i in range(batch_size):
            traj = trajectories[int(sorted_inds[batch_inds[i]])]
            si = random.randint(0, traj['rewards'].shape[0] - 1)
            s_all = traj['observations']
            s.append(traj['observations'][si:si + max_len].reshape(1, -1, state_dim))
            tlen = s[-1].shape[1]
            target_rtg.append(np.asarray(traj['return_reals'][si:si + max_len]).reshape(1, -1, 1))
            closest_indices = np.asarray(traj['closest_indices'][si:si + max_len]).reshape(-1,num_subgoal)
            s[-1] = np.concatenate([np.zeros((1, max_len - tlen, state_dim)), s[-1]], axis=1)
            target_rtg[-1] = np.concatenate([np.zeros((1, max_len - tlen, 1)), target_rtg[-1]], axis=1) / scale
            s[-1] = (s[-1] - state_mean) / state_std
            closest_indices = np.concatenate([closest_indices,np.full((max_len - tlen,num_subgoal),-1)],axis=0)
            result = np.zeros((closest_indices.shape[0], closest_indices.shape[1], state_dim))
            valid_indices = closest_indices != -1
            result[valid_indices]=s_all[closest_indices[valid_indices]]
            subgoal.append(result.reshape(1,max_len,num_subgoal,state_dim))
        s = torch.from_numpy(np.concatenate(s, axis=0)).to(dtype=torch.float32, device=device)
        all_rtg = torch.tensor(target_rtg)
        coeffs = torch.tensor(generate_arithmetic_sequence(num_subgoal) ).view(1, 1, num_subgoal)
        all_rtg = (all_rtg * coeffs).reshape(batch_size,-1,num_subgoal,1)
        target_rtg = torch.Tensor(target_rtg).to(dtype=torch.float32, device=device)
        all_rtg = torch.tensor(all_rtg,dtype=torch.float32, device=device)
        # plan_masks = (plan_masks).to(device=device)
        subgoal = torch.from_numpy(np.concatenate(subgoal, axis=0)).to(dtype=torch.float32, device=device)
        tgt_mask = nn.Transformer.generate_square_subsequent_mask(num_subgoal*2).to(dtype=torch.float32, device=device)
        return s,target_rtg,subgoal,tgt_mask,all_rtg
    def get_batch(batch_size=256,num_subgoal=5,max_len=K):
        batch_inds = np.random.choice(
            np.arange(num_trajectories),
            size=batch_size,
            replace=True,
            p=p_sample,  # reweights so we sample according to timesteps
        )
        s, a, r, d, rtg, timesteps, mask,target_rtg,subgoal=  [], [], [], [], [],  [], [] ,[],[]
        for i in range(batch_size):
            traj = trajectories[int(sorted_inds[batch_inds[i]])]
            si = random.randint(0,traj['rewards'].shape[0]-1)
            s.append(traj['observations'][si:si + max_len].reshape(1, -1, state_dim))
            a.append(traj['actions'][si:si + max_len].reshape(1, -1, act_dim))
            r.append(traj['rewards'][si:si + max_len].reshape(1, -1, 1))
            if 'terminals' in traj:
                d.append(traj['terminals'][si:si + max_len].reshape(1, -1))
            else:
                d.append(traj['dones'][si:si + max_len].reshape(1, -1))
            timesteps.append(np.arange(si, si + s[-1].shape[1]).reshape(1, -1))
            timesteps[-1][timesteps[-1] >= max_ep_len] = max_ep_len-1  # padding cutoff
            rtg.append(discount_cumsum(traj['rewards'][si:], gamma=1.)[:s[-1].shape[1]].reshape(1, -1, 1))
            target_rtg.append(np.asarray(traj['return_reals'][si:si + max_len]).reshape(1, -1, 1))
            s_all = traj['observations']
            tlen = s[-1].shape[1]
            closest_indices = np.asarray(traj['closest_indices'][si:si + max_len]).reshape(-1,num_subgoal)
            closest_indices = np.concatenate([closest_indices,np.full((max_len - tlen,num_subgoal),-1)],axis=0)
            result = np.zeros((closest_indices.shape[0], closest_indices.shape[1], state_dim))
            valid_indices = closest_indices != -1
            result[valid_indices]=s_all[closest_indices[valid_indices]]
            subgoal.append(result.reshape(1,max_len,num_subgoal,state_dim))
            # padding and state + reward normalization
            s[-1] = np.concatenate([np.zeros((1, max_len - tlen, state_dim)), s[-1]], axis=1)
            s[-1] = (s[-1] - state_mean) / state_std
            a[-1] = np.concatenate([np.ones((1, max_len - tlen, act_dim)) * -10., a[-1]], axis=1)
            r[-1] = np.concatenate([np.zeros((1, max_len - tlen, 1)), r[-1]], axis=1)
            d[-1] = np.concatenate([np.ones((1, max_len - tlen)) * 2, d[-1]], axis=1)
            rtg[-1] = np.concatenate([np.zeros((1, max_len - tlen, 1)), rtg[-1]], axis=1) / scale
            timesteps[-1] = np.concatenate([np.zeros((1, max_len - tlen)), timesteps[-1]], axis=1)
            target_rtg[-1] = np.concatenate([np.zeros((1, max_len - tlen, 1)), target_rtg[-1]], axis=1) / scale
            # plan_masks.append(plan_mask.reshape(1,max_len,-1))
            mask.append(np.concatenate([np.zeros((1, max_len - tlen)), np.ones((1, tlen))], axis=1))
        s = torch.from_numpy(np.concatenate(s, axis=0)).to(dtype=torch.float32, device=device)
        a = torch.from_numpy(np.concatenate(a, axis=0)).to(dtype=torch.float32, device=device)
        r = torch.from_numpy(np.concatenate(r, axis=0)).to(dtype=torch.float32, device=device)
        d = torch.from_numpy(np.concatenate(d, axis=0)).to(dtype=torch.long, device=device)
        rtg = torch.from_numpy(np.concatenate(rtg, axis=0)).to(dtype=torch.float32, device=device)
        timesteps = torch.from_numpy(np.concatenate(timesteps, axis=0)).to(dtype=torch.long, device=device)
        mask = torch.from_numpy(np.concatenate(mask, axis=0)).to(device=device)
        subgoal = torch.from_numpy(np.concatenate(subgoal, axis=0)).to(dtype=torch.float32, device=device)
        target_rtg = torch.Tensor(target_rtg).to(dtype=torch.float32, device=device)
        # plan_masks = (plan_masks).to(device=device)
        return s, a, r, d, rtg, timesteps, mask ,target_rtg,subgoal
    def eval_episodes(target_rew):
        def fn(model):
            returns, lengths,norm_scores = [], [], []
            for _ in range(num_eval_episodes):
                with torch.no_grad():
                    if model_type == 'dt':
                        ret, length,norm_score = evaluate_episode_rtg(
                            env,
                            state_dim,
                            act_dim,
                            future_dim,
                            model,
                            max_ep_len=max_ep_len,
                            scale=scale,
                            target_return=target_rew/scale,
                            mode=mode,
                            state_mean=state_mean,
                            state_std=state_std,
                            device=device,
                        )
                    else:
                        ret, length = evaluate_episode(
                            env,
                            state_dim,
                            act_dim,
                            model,
                            max_ep_len=max_ep_len,
                            target_return=target_rew/scale,
                            mode=mode,
                            state_mean=state_mean,
                            state_std=state_std,
                            device=device,
                        )

                returns.append(ret)
                lengths.append(length)
                norm_scores.append(norm_score)
            return {
                f'target_{target_rew}_return_mean': np.mean(returns),
                f'target_{target_rew}_return_std': np.std(returns),
                f'target_{target_rew}_length_mean': np.mean(lengths),
                f'target_{target_rew}_length_std': np.std(lengths),
                f'target_{target_rew}_norm_scores_std': np.std(norm_scores),
                f'target_{target_rew}_norm_scores_mean': np.mean(norm_scores),
            }
        return fn

    if model_type == 'dt':
        model = Cross_attetion_DecisionTransformer(
            device =device,
            state_dim=state_dim,
            act_dim=act_dim,
            max_length=K,
            max_ep_len=max_ep_len,
            nom_subgoal = num_subgoal,
            hidden_size=variant['embed_dim'],
            n_layer=variant['n_layer'],
            n_head=variant['n_head'],
            n_inner=4*variant['embed_dim'],
            activation_function=variant['activation_function'],
            n_positions=1024,
            resid_pdrop=variant['dropout'],
            attn_pdrop=variant['dropout'],
            futures_dim=variant['futures_dim'],
            d_model = variant['d_model'],
            nhead = variant['nhead'], 
            num_layers = variant['num_layers']
        )
    elif model_type == 'bc':
        model = MLPBCModel(
            state_dim=state_dim,
            act_dim=act_dim,
            max_length=K,
            hidden_size=variant['embed_dim'],
            n_layer=variant['n_layer'],
        )
    else:
        raise NotImplementedError

    model = model.to(device=device)

    warmup_steps = variant['warmup_steps']
    optimizer_pre = torch.optim.AdamW(
        model.parameters(),
        lr=variant['learning_rate_pre'],
        weight_decay=variant['weight_decay'],
    )
    optimizer = torch.optim.AdamW(
        model.parameters(),
        lr=variant['learning_rate'],
        weight_decay=variant['weight_decay'],
    )
    scheduler = torch.optim.lr_scheduler.LambdaLR(
        optimizer,
        lambda steps: min((steps+1)/warmup_steps, 1)
    )

    if model_type == 'dt':
        trainer = Cross_attetion_Trainer(
            model=model,
            optimizer=optimizer,
            optimizer_pre = optimizer_pre,
            batch_size=batch_size,
            num_subgoal = num_subgoal,
            scheduler=scheduler,
            device =device,
            loss_fn=lambda s_hat, a_hat, r_hat, s, a, r: torch.mean((a_hat - a)**2),
            eval_fns=[eval_episodes(tar) for tar in env_targets],
            get_batch=get_batch,
            get_batch_new = get_batch_new
        )
    elif model_type == 'bc':
        trainer = ActTrainer(
            model=model,
            optimizer=optimizer,
            batch_size=batch_size,
            get_batch=get_batch,
            scheduler=scheduler,
            loss_fn=lambda s_hat, a_hat, r_hat, s, a, r: torch.mean((a_hat - a)**2),
            eval_fns=[eval_episodes(tar) for tar in env_targets],
        )

    if log_to_swanlab:
        run = swanlab.init(
                workspace="PlanDT",
                project=f'{env_name}-{dataset}',
                experiment_name=f'{env_name}-{dataset}-plan-cross-attention-subgoal({num_subgoal})-temperatrue({temperature})'
        )
    for iter in range(variant['max_iters']):
        outputs = trainer.train_encoder_decoder(num_steps=variant['num_steps_per_iter_pre'],temprature =variant['temperature'])
        if log_to_swanlab:
            run.log(outputs)
    for iter in range(variant['max_iters']):
        outputs = trainer.train_iteration(num_steps=variant['num_steps_per_iter'])
        if log_to_swanlab:
            run.log(outputs)
    print("Hyperparameters used in this experiment:")
    for arg, value in vars(args).items():
            print(f"{arg}: {value}")

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--env', type=str, default='maze2d')
    parser.add_argument('--dataset', type=str, default='medium-dense')  # medium, medium-replay, medium-expert, expert
    parser.add_argument('--mode', type=str, default='normal')  # normal for standard setting, delayed for sparse
    parser.add_argument('--K', type=int, default=20)
    parser.add_argument('--pct_traj', type=float, default=1.)
    parser.add_argument('--batch_size', type=int, default=128)
    parser.add_argument('--model_type', type=str, default='dt')
    parser.add_argument('--embed_dim', type=int, default=128)
    parser.add_argument('--n_layer', type=int, default=2)
    parser.add_argument('--n_head', type=int, default=16)
    parser.add_argument('--activation_function', type=str, default='relu')
    parser.add_argument('--dropout', type=float, default=0.15)
    parser.add_argument('--learning_rate', '-lr', type=float, default=1e-4)
    parser.add_argument('--learning_rate_pre', type=float, default=1e-4)
    parser.add_argument('--weight_decay', '-wd', type=float, default=1e-4)
    parser.add_argument('--warmup_steps', type=int, default=10000)#10000
    parser.add_argument('--num_eval_episodes', type=int, default=10)
    parser.add_argument('--max_iters', type=int, default=20)
    parser.add_argument('--num_steps_per_iter_pre', type=int, default=2000)
    parser.add_argument('--num_steps_per_iter', type=int, default=5000)#10000
    parser.add_argument('--device', type=str, default='cuda:1')
    parser.add_argument('--temperature', type=float, default=0.1)
    parser.add_argument('--num_subgoal', type=int ,default=3)
    parser.add_argument('--futures_dim', type=int ,default=16)
    parser.add_argument('--d_model', type=int ,default=128)
    parser.add_argument('--num_layers', type=int, default=24)
    parser.add_argument('--nhead', type=int ,default=16)
    parser.add_argument('--log_to_swanlab', '-w', type=bool, default=True)
    args = parser.parse_args()

    experiment(f'plan-cross-attetion-subgoal(10)-drop(0.15)-50(2000-10000)-learning_rate(1e-4)-256', variant=vars(args))