import sys

import argparse

import os

import json

import csv

import numpy as np

from utils.filemanager import get_points_paths

from utils.logger import logger, pprint

from utils.landmarks import get_landmarks_from_txt, write_landmarks_to_list

from utils.metrics import compute_TRE

if __name__ == "__main__":

    # optional arguments from the command line 

    parser = argparse.ArgumentParser()

    parser.add_argument('--experiment_name', type=str, default='elastix_01', help='experiment name')

    parser.add_argument('--reg_params_key', type=str, default='Parameter.affine+Parameter.bsplines', help='registration parameters key generated by create_script.py')

    parser.add_argument('--output_path', type=str, default='output', help='root dir for output scripts')

    parser.add_argument("--generate_report", action='store_true', help='if True, an evaluation report .txt file will be generated. If not, only the transformed keypoints txt file will be generated for each test sample.')

    parser.add_argument('--dataset_path', type=str, default='dataset/train', help='root dir for nifti data to get the gt exhale landmarks')

    # parse the arguments

    args = parser.parse_args()

    # create experiment search path

    # points is the folder where the transformed points are saved using transformix

    args.exp_points_output = os.path.join(args.output_path, args.experiment_name, args.reg_params_key, 'points')

    # get a list of all the transformed keypoints files

    transformed_points = get_points_paths(args.exp_points_output, "outputpoints", num_occurrences=2)

    if len(transformed_points) == 0:

        logger.error(f"No transformed points found in {args.exp_points_output} directory.")

        sys.exit(1)

    # check if generate_report is True

    if args.generate_report:

        gt_points = get_points_paths(args.dataset_path, "_300_eBH_xyz_r1", num_occurrences=1)

        if len(gt_points) == 0:

            logger.error(f"No gt points found in {args.dataset_path} directory.")

            sys.exit(1)

        # Create a list to store the TRE results

        tre_results = []

    else:

        gt_points = [0 for _ in range(len(transformed_points))] # the list has to have values for the zip(*) to return the values inside

    logger.info(f"Found {len(transformed_points)} transformed points files for subjects ({[subject.split('/')[-2] for subject in transformed_points]})")

    # extract the transformed points from the transformed_points transformix files and save them in a separate file

    for transformed_points_file, gt_point in zip(transformed_points, gt_points):

        print(f"Processing {transformed_points_file}...")

        # get the transformed points

        transformed_landmarks = get_landmarks_from_txt(transformed_points_file, search_key='OutputIndexFixed')

        # the transformed points has to be 300

        assert len(transformed_landmarks) == 300, f"Transformed points file {transformed_points_file} has {len(transformed_landmarks)} points instead of 300."

        # write the transformed points to a file 

        # the points are written inside the same directory as the transformed_points_file

        output_landmarks_path = os.path.join(transformed_points_file.replace('outputpoints.txt', ''), 'outputpoints_transformed.txt')

        write_landmarks_to_list(transformed_landmarks, output_landmarks_path)

        # generate the evaluation report if args.generate_report is True, this is when we have the ground truth exhale files

        if args.generate_report:

            sample_name = gt_point.split('/')[-1].split('_')[0] #copd1, copd2, ...

            # load the dataset dictionary, we remove the last path element because we want to get the description.json file

            with open(os.path.join(args.dataset_path.replace("train", "", 1).replace("test", "", 1),'description.json'), 'r') as json_file:

                dictionary = json.loads(json_file.read())

            file_information = dictionary[args.dataset_path.replace('\\', '/').split("/")[-1]][sample_name]

            print(file_information)

            TRE_mean, TRE_std = compute_TRE(output_landmarks_path, gt_point, tuple(file_information['voxel_dim']))

            print("TRE (After Registration):- ", f"(Mean TRE: {TRE_mean})", f"(STD TRE: {TRE_std}). \n")

            # Append TRE results to the list

            tre_results.append({'sample_name': sample_name, 'TRE_mean': TRE_mean, 'TRE_std': TRE_std})

    # generate the evaluation report if args.generate_report is True, this is when we have the ground truth exhale files

    if args.generate_report:

        # write the TRE results to a csv file for each sample

        output_csv_path = os.path.join(args.exp_points_output, 'TRE_sample_results.csv')

        with open(output_csv_path, 'w', newline='') as csv_file:

            fieldnames = ['sample_name', 'TRE_mean', 'TRE_std']

            writer = csv.DictWriter(csv_file, fieldnames=fieldnames)

            # Write the header

            writer.writeheader()

            # Write the data

            for result in tre_results:

                writer.writerow(result)

        # write the overall mean results

        TRE_mean_list = [result['TRE_mean'] for result in tre_results]

        TRE_std_list = [result['TRE_std'] for result in tre_results]

        output_csv_path = os.path.join(args.exp_points_output, 'TRE_overall_results.csv')

        with open(output_csv_path, 'w', newline='') as csv_file:

            fieldnames = ['Overall mean (TRE_mean)', 'Overall mean (TRE_std)']

            writer = csv.DictWriter(csv_file, fieldnames=fieldnames)

            # Write the header

            writer.writeheader()

            # Write the data

            writer.writerow({'Overall mean (TRE_mean)': np.mean(TRE_mean_list), 'Overall mean (TRE_std)': np.mean(TRE_std_list)})