import os
import sys
import time
import torch
import argparse
import torchvision
import numpy as np
import torch.nn as nn
import torch.nn.functional as F

from logger_utils import CSVWriter
from model_main import CRNN, STN_CRNN
from utils import ctc_decode, compute_wer_and_cer_for_sample
from dataset import HWRecogIAMDataset, split_dataset, get_dataloader_for_testing


def test(hw_model, test_loader, device, list_test_files, which_ctc_decoder="beam_search", save_prediction_stats=False):
    """
    ---------
    Arguments
    ---------
    hw_model : object
        handwriting recognition model object
    test_loader : object
        dataset loader object
    device : str
        device to be used for running the evaluation
    list_test_files : list
        list of all the test files
    which_ctc_decoder : str
        string indicating which ctc decoder to use
    save_prediction_stats : bool
        whether to save prediction stats
    """
    hw_model.eval()
    num_test_samples = len(test_loader.dataset)
    num_test_batches = len(test_loader)

    count = 0
    list_test_cers, list_test_wers = [], []

    if save_prediction_stats:
        csv_writer = CSVWriter(
            file_name="pred_stats.csv",
            column_names=["file_name", "num_chars", "num_words", "cer", "wer"]
        )

    with torch.no_grad():
        for images, labels, length_labels in test_loader:
            count += 1
            images = images.to(device, dtype=torch.float)
            log_probs = hw_model(images)
            pred_labels = ctc_decode(log_probs, which_ctc_decoder=which_ctc_decoder)
            labels = labels.cpu().numpy().tolist()

            str_label = [HWRecogIAMDataset.LABEL_2_CHAR[i] for i in labels]
            str_label = "".join(str_label)
            str_pred = [HWRecogIAMDataset.LABEL_2_CHAR[i] for i in pred_labels[0]]
            str_pred = "".join(str_pred)

            cer_sample, wer_sample = compute_wer_and_cer_for_sample(str_pred, str_label)
            list_test_cers.append(cer_sample)
            list_test_wers.append(wer_sample)

            print(f"progress: {count}/{num_test_samples}, test file: {list_test_files[count-1]}")
            print(f"{str_label} - label")
            print(f"{str_pred} - prediction")
            print(f"cer: {cer_sample:.3f}, wer: {wer_sample:.3f}\n")

            if save_prediction_stats:
                csv_writer.write_row([
                    list_test_files[count-1],
                    len(str_label),
                    len(str_label.split(" ")),
                    cer_sample,
                    wer_sample,
                ])
    list_test_cers = np.array(list_test_cers)
    list_test_wers = np.array(list_test_wers)
    mean_test_cer = np.mean(list_test_cers)
    mean_test_wer = np.mean(list_test_wers)
    print(f"test set - mean cer: {mean_test_cer:.3f}, mean wer: {mean_test_wer:.3f}\n")

    if save_prediction_stats:
        csv_writer.close()
    return

def test_hw_recognizer(FLAGS):
    file_txt_labels = os.path.join(FLAGS.dir_dataset, "iam_lines_gt.txt")
    dir_images = os.path.join(FLAGS.dir_dataset, "img")
    os.environ["CUDA_VISIBLE_DEVICES"] = "0,1"

    # choose a device for testing
    if torch.cuda.is_available():
        device = torch.device("cuda")
    else:
        device = torch.device("cpu")

    # get the internal test set files
    test_x, test_y = split_dataset(file_txt_labels, for_train=False)
    num_test_samples = len(test_x)
    # get the internal test set dataloader
    test_loader = get_dataloader_for_testing(
        test_x, test_y,
        dir_images=dir_images, image_height=FLAGS.image_height, image_width=FLAGS.image_width,
    )

    num_classes = len(HWRecogIAMDataset.LABEL_2_CHAR) + 1
    print(f"task - handwriting recognition")
    print(f"model: {FLAGS.which_hw_model}, ctc decoder: {FLAGS.which_ctc_decoder}")
    print(f"image height: {FLAGS.image_height}, image width: {FLAGS.image_width}")
    print(f"num test samples: {num_test_samples}")

    # load the right model
    if FLAGS.which_hw_model == "crnn":
        hw_model = CRNN(num_classes, FLAGS.image_height)
    elif FLAGS.which_hw_model == "stn_crnn":
        hw_model = STN_CRNN(num_classes, FLAGS.image_height, FLAGS.image_width)
    else:
        print(f"unidentified option : {FLAGS.which_hw_model}")
        sys.exit(0)
    hw_model.to(device)
    hw_model.load_state_dict(torch.load(FLAGS.file_model))

    # start testing of the model on the internal set
    print(f"testing of handwriting recognition model {FLAGS.which_hw_model} started\n")
    test(hw_model, test_loader, device, test_x, FLAGS.which_ctc_decoder, bool(FLAGS.save_prediction_stats))
    print(f"testing handwriting recognition model completed!!!!")
    return

def main():
    image_height = 32
    image_width = 768
    which_hw_model = "crnn"
    dir_dataset = "/home/abhishek/Desktop/RUG/hw_recognition/IAM-data/"
    file_model = "model_crnn/crnn_H_32_W_768_E_177.pth"
    which_ctc_decoder = "beam_search"
    save_prediction_stats = 0

    parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )

    parser.add_argument("--image_height", default=image_height,
        type=int, help="image height to be used to predict with the model")
    parser.add_argument("--image_width", default=image_width,
        type=int, help="image width to be used to predict with the model")
    parser.add_argument("--dir_dataset", default=dir_dataset,
        type=str, help="full directory path to the dataset")
    parser.add_argument("--which_hw_model", default=which_hw_model,
        type=str, choices=["crnn", "stn_crnn"], help="which model to be used for prediction")
    parser.add_argument("--which_ctc_decoder", default=which_ctc_decoder,
        type=str, choices=["beam_search", "greedy"], help="which ctc decoder to use")
    parser.add_argument("--file_model", default=file_model,
        type=str, help="full path to trained model file (.pth)")
    parser.add_argument("--save_prediction_stats", default=save_prediction_stats,
        type=int, choices=[0, 1], help="save prediction stats (1 - yes, 0 - no)")

    FLAGS, unparsed = parser.parse_known_args()
    test_hw_recognizer(FLAGS)
    return

if __name__ == "__main__":
    main()