import argparse
import os
import jax
import flax
import copy
import json
import optax
import torch
import wandb
import numpy as np
from tqdm import tqdm
import jax.numpy as jnp
from flax import linen as nn
from flax.jax_utils import replicate, unreplicate
from flax.training import checkpoints, train_state
from flax.core.frozen_dict import freeze, unfreeze
from flax.traverse_util import flatten_dict, unflatten_dict
from transformers.models.vit.modeling_flax_vit import ViTConfig, FlaxViTForImageClassification
from flax.training.common_utils import get_metrics, onehot, shard
from lmc_model import  LMCFlaxViTForImageClassification, print_model, print_model_with_prefix
from datasets import build_dataset
import multiprocessing as mp
from jax import debug
from pprint import pprint
from typing import Any, Dict, List
import shutil
mp.set_start_method("spawn", force=True)
os.environ["WANDB_API_KEY"] = "fc72050bcc0dc7f7502b5416938f8bd0c4b30fc7"
def imagenet_data_loader(args):
    dataset_train, args.nb_classes = build_dataset(is_train=True, args=args)
    dataset_val, _ = build_dataset(is_train=False, args=args)
    sampler_train = torch.utils.data.RandomSampler(dataset_train)
    sampler_val = torch.utils.data.SequentialSampler(dataset_val)
    data_loader_train = torch.utils.data.DataLoader(
        dataset_train, sampler=sampler_train, batch_size=args.batch_size,
        num_workers=args.num_workers, pin_memory=args.pin_mem, drop_last=True
    )
    data_loader_val = torch.utils.data.DataLoader(
        dataset_val, sampler=sampler_val, batch_size=args.batch_size,
        num_workers=args.num_workers, pin_memory=args.pin_mem, drop_last=False
    )
    return data_loader_train, data_loader_val
def main(args: argparse.Namespace):
    save_path = "/mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256-shared1-routed0-topk0"
    model = LMCFlaxViTForImageClassification.from_pretrained("/mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256-shared1-routed0-topk0/temp_1291032")
    step = 1291032
    num_global_steps = 5004*args.epochs 
    num_warmup_steps = 5004*args.warmup_epochs
    lr_schedule =optax.warmup_cosine_decay_schedule(
        init_value=args.warmup_lr,
        peak_value=args.lr,
        warmup_steps=num_warmup_steps,
        decay_steps=num_global_steps,
        end_value=args.min_lr,
    )
    tx = optax.adamw(
        learning_rate=lr_schedule,
        b1=args.adamw_beta1,
        b2=args.adamw_beta2,
        eps=args.adamw_eps,
        weight_decay=args.weight_decay,
    )
    state = train_state.TrainState.create(apply_fn=model.__call__, params=model.params, tx=tx)
    state = state.replace(step=step)
    checkpoints.save_checkpoint(ckpt_dir=save_path,target=state,step=state.step,prefix="last_",keep=1)                    


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Fine-tune ViT with MoE on Imagenet")
    parser.add_argument("--epochs", type=int, default=30)
    parser.add_argument("--batch-size", type=int, default=64)
    parser.add_argument("--lr", type=float, default=5e-4)
    parser.add_argument("--weight-decay", type=float, default=0.01)
    parser.add_argument('--sched', default='cosine', type=str, metavar='SCHEDULER')
    parser.add_argument('--warmup-lr', type=float, default=1e-6, metavar='LR', help='warmup learning rate (default: 1e-6)')
    parser.add_argument('--warmup-epochs', type=int, default=3, metavar='N',help='epochs to warmup LR, if scheduler supports')
    parser.add_argument('--min-lr', type=float, default=1e-5, metavar='LR',help='lower lr bound for cyclic schedulers that hit 0 (1e-5)')
    parser.add_argument("--adamw-beta1", type=float, default=0.9)
    parser.add_argument("--adamw-beta2", type=float, default=0.999)    
    parser.add_argument("--adamw-eps", type=float, default=1e-8)
    parser.add_argument("--patience", type=int, default=10, help="Early stopping patience")
    args = parser.parse_args()
    main(args)