import argparse
import os
import pandas as pd

from torch.utils.data import DataLoader

from model_zoo import get_model
from dataset_zoo import get_dataset
from dataset_zoo.perturbations import get_text_perturb_fn, get_image_perturb_fn
from misc import seed_all, _default_collate, save_scores


def config():
    parser = argparse.ArgumentParser()
    parser.add_argument("--device", default="cuda", type=str)
    parser.add_argument("--batch-size", default=32, type=int)
    parser.add_argument("--num_workers", default=4, type=int)
    parser.add_argument("--model-name", default="openai-clip:ViT-B/32", type=str)
    parser.add_argument(
        "--dataset",
        default="COCO_Retrieval",
        type=str,
        choices=["COCO_Retrieval", "Flickr30k_Retrieval"],
    )
    parser.add_argument("--seed", default=1, type=int)
    parser.add_argument(
        "--text-perturb-fn",
        default=None,
        type=str,
        help="Perturbation function to apply to the text.",
    )
    parser.add_argument(
        "--image-perturb-fn",
        default=None,
        type=str,
        help="Perturbation function to apply to the images.",
    )

    parser.add_argument(
        "--download",
        action="store_true",
        help="Whether to download the dataset if it doesn't exist. (Default: False)",
    )
    parser.add_argument(
        "--save-scores",
        action="store_true",
        help="Whether to save the scores for the retrieval to analyze later.",
    )
    parser.add_argument("--output-dir", default="./outputs", type=str)
    return parser.parse_args()


def main(args):
    seed_all(args.seed)

    model, image_preprocess = get_model(args.model_name, args.device)
    text_perturb_fn = get_text_perturb_fn(args.text_perturb_fn)
    image_perturb_fn = get_image_perturb_fn(args.image_perturb_fn)

    dataset = get_dataset(
        args.dataset,
        image_preprocess=image_preprocess,
        text_perturb_fn=text_perturb_fn,
        image_perturb_fn=image_perturb_fn,
        download=args.download,
    )
    # For some models we just pass the PIL images, so we'll need to handle them in the collate_fn.
    collate_fn = _default_collate if image_preprocess is None else None

    loader = DataLoader(
        dataset,
        batch_size=args.batch_size,
        shuffle=False,
        num_workers=args.num_workers,
        collate_fn=collate_fn,
    )

    if text_perturb_fn is not None:
        loader.dataset.text = [text_perturb_fn(t) for t in loader.dataset.text]

    scores = model.get_retrieval_scores_dataset(loader)
    result_records = dataset.evaluate_scores(scores)

    for record in result_records:
        record.update(
            {
                "Model": args.model_name,
                "Dataset": args.dataset,
                "Text Perturbation Strategy": args.text_perturb_fn,
                "Seed": args.seed,
                "Image Perturbation Strategy": args.image_perturb_fn,
            }
        )

    df = pd.DataFrame(result_records)
    output_file = os.path.join(args.output_dir, f"{args.dataset}.csv")
    print(f"Saving results to {output_file}")
    if os.path.exists(output_file):
        all_df = pd.read_csv(output_file, index_col=0)
        all_df = pd.concat([all_df, df])
        all_df.to_csv(output_file)

    else:
        df.to_csv(output_file)

    if args.save_scores:
        save_scores(scores, args)


if __name__ == "__main__":
    args = config()
    main(args)
