import argparse
import time
import os
import numpy as np
from tsn.models.TSN import TSN
from tsn.utils.checkpoint import load_checkpoint
import oneflow as flow
from tsn.datasets.transform import *
from tsn.datasets.dataset import TSNDataSet

import warnings

warnings.filterwarnings("ignore", category=UserWarning)


def parse_args():
    parser = argparse.ArgumentParser(description="Test an action recognizer")
    parser.add_argument("--test_mode", default=False, help="train or test mode")
    parser.add_argument("-load_checkpoint", default="", help="checkpoint file")
    parser.add_argument(
        "--data_dir", default="./data", help="data file path",
    )
    parser.add_argument(
        "--pretrained",
        default="./resnet50_imagenet_pretrain_model",
        help="test videos prefix path",
    )
    parser.add_argument(
        "--input_mean", default=[0.485, 0.456, 0.406], help="data normalization value"
    )
    parser.add_argument(
        "--input_std", default=[0.229, 0.224, 0.225], help="data normalization value"
    )
    parser.add_argument("--epochs", default=100, help="max epochs for training")
    parser.add_argument("--lr_steps", default=[40, 80], help="lr step")
    parser.add_argument("--batch_size", default=4, help="imgs per gpu")
    parser.add_argument(
        "--save_checkpoint_path", default="./res", help="imgs per gpu",
    )
    parser.add_argument("--lr", default=0.0025, help="learning_rate")
    parser.add_argument("--mom", default=0.9, help="momentum")
    parser.add_argument("--weight_decay", default=0.000, help="weight decay")
    # for spatial_temporal module
    parser.add_argument(
        "--spatial_type", default="avg", help="data normalization value"
    )
    parser.add_argument("--spatial_size", default=7, help="data normalization value")
    # for segmental consensus
    parser.add_argument(
        "--consensus_type", default="avg", help="data normalization value"
    )
    # for class head
    parser.add_argument(
        "--spatial_feature_size", default=1, help="data normalization value"
    )
    parser.add_argument("--dropout_ratio", default=0.4, help="data normalization value")
    parser.add_argument("--num_classes", default=400, help="data normalization value")
    parser.add_argument("--out", help="output result file", default="default.pkl")
    parser.add_argument(
        "--use_softmax", action="store_true", help="whether to use softmax score"
    )
    # only for TSN3D
    parser.add_argument(
        "--fcn_testing", action="store_true", help="use fcn testing for 3D convnet"
    )
    parser.add_argument("--local_rank", type=int, default=0)

    args = parser.parse_args()
    return args


def top_k_hit(score, lb_set, k=3):
    idx = np.argsort(score)[-k:]
    return len(lb_set.intersection(idx)) > 0, 1


def top_k_accuracy(scores, labels, k=(1,)):
    res = []
    for kk in k:
        hits = []
        for x, y in zip(scores, labels):
            y = [y] if isinstance(y, int) else y
            hits.append(top_k_hit(x, set(y), k=kk)[0])
        res.append(np.mean(hits))
    return res


def adjust_learning_rate(optimizer, epoch, lr_steps):
    """Sets the learning rate to the initial LR decayed by lr steps"""
    decay = 0.1 ** (sum(epoch >= np.array(lr_steps)))
    lr = args.lr * decay

    for param_group in optimizer.param_groups:
        param_group["lr"] = lr


def main():

    global args
    args = parse_args()

    model = TSN(
        args.spatial_feature_size, args.dropout_ratio, args.num_classes, args.pretrained
    )

    train_ann_file = (
        args.data_dir + "/kinetics400/kinetics400_train_list_videos_sub_frames.txt"
    )
    train_prefix = args.data_dir + "/kinetics400/rawframes_train"
    val_ann_file = (
        args.data_dir + "/kinetics400/kinetics400_val_list_videos_sub_frames.txt"
    )
    val_prefix = args.data_dir + "/kinetics400/rawframes_val"

    train_dataset = TSNDataSet(
        "",
        train_ann_file,
        num_segments=3,
        video_dir=train_prefix,
        new_length=1,
        image_tmpl="img_{:05d}.jpg",
        flip=GroupMultiScaleCrop(224, [1, 0.875, 0.75, 0.66]),
        crop=GroupRandomHorizontalFlip(is_flow=False),
        stack=Stack(roll=False),
        Normalize=GroupNormalize(args.input_mean, args.input_std),
        batch_size=args.batch_size,
    )

    val_dataset = TSNDataSet(
        "",
        val_ann_file,
        num_segments=3,
        video_dir=val_prefix,
        new_length=1,
        image_tmpl="img_{:05d}.jpg",
        test_mode=True,
        sample=GroupOverSample(224, 256),
        stack=Stack(roll=False),
        Normalize=GroupNormalize(args.input_mean, args.input_std),
        batch_size=args.batch_size,
    )

    if args.load_checkpoint != "":
        model.load_state_dict(flow.load(args.load_checkpoint))

    of_corss_entropy = flow.nn.CrossEntropyLoss()
    of_corss_entropy.to("cuda")

    model.to("cuda")

    optimizer = flow.optim.SGD(model.parameters(), lr=args.lr, momentum=args.mom)

    losses = []
    print_interval = 50
    val_interval = 1
    for epoch in range(args.epochs):
        adjust_learning_rate(optimizer, epoch, args.lr_steps)
        model.train()

        for idx in range(int(len(train_dataset) / args.batch_size)):
            data, label = train_dataset[idx]
            num_modalities = 1
            img_group = data.reshape(
                args.batch_size,
                int(data.shape[0] / args.batch_size),
                data.shape[1],
                data.shape[2],
                data.shape[3],
            )
            start_t = time.time()
            img_group = flow.Tensor(img_group, device=flow.device("cuda"))
            gt_label = flow.Tensor(label, device=flow.device("cuda"))
            gt_label = gt_label.squeeze().to(dtype=flow.int32)

            predeicts = model(num_modalities, gt_label, img_group, True)
            loss = of_corss_entropy(predeicts, gt_label)
            loss.backward()
            optimizer.step()
            optimizer.zero_grad()
            end_t = time.time()

            if idx % print_interval == 0:
                l = loss.numpy()
                losses.append(l)
                tmp_lr = optimizer.param_groups[0]["lr"]
                print(
                    "epoch {} train iter {} oneflow loss {}, lr: {}, train time : {}".format(
                        epoch, idx, l, tmp_lr, end_t - start_t
                    )
                )

        if epoch % val_interval == 0:
            model.eval()
            results = []

            for idx in range(int(len(val_dataset) / args.batch_size)):
                data, label = val_dataset[idx]
                num_modalities = 1
                img_meta = label
                img_group = data.reshape(
                    args.batch_size,
                    int(data.shape[0] / args.batch_size),
                    data.shape[1],
                    data.shape[2],
                    data.shape[3],
                )
                img_group = flow.Tensor(img_group, device=flow.device("cuda"))
                with flow.no_grad():
                    result = model(num_modalities, img_meta, img_group)
                results.append(result)

            outputs = []
            for res in results:
                for idx in range(0, args.batch_size):
                    outputs.append(res[idx])

            gt_labels = []
            for i in range(len(val_dataset)):
                ann = val_dataset.get_ann_info(i)
                gt_labels.append(ann["label"])

            top1, top5 = top_k_accuracy(outputs, gt_labels, k=(1, 5))
            print(
                "epoch %d, oneflow top1 val acc: %f, top5 val acc: %f"
                % (epoch, top1 * 100, top5 * 100)
            )
            flow.save(
                model.state_dict(),
                os.path.join(
                    args.save_checkpoint_path,
                    "epoch_%d_val_acc_%f_%f" % (epoch, top1, top5),
                ),
            )

    writer = open("of_losses.txt", "w")
    for o in losses:
        writer.write("%f\n" % o)
    writer.close()


if __name__ == "__main__":
    main()
