from tqdm import tqdm

import aix
from aix import AIX_DATASETS, AIX_EVALS, MATURE, IMMATURE
from aix.np_metrics import hard_dice, dice, rel_size_diff
import cv2
import pandas as pd
import json

all_measures = {
        "dice": dice,
        "hard_dice": hard_dice,
        "rel_size_diff": rel_size_diff
    }

OOCYTE = "oocyte"
REFERENCE = "reference"
COMPARED = "compared"


def load_csv(filename):
    csv_file_path = AIX_EVALS / filename
    if csv_file_path.is_file():
        print("Loaded")
        return pd.read_csv(csv_file_path, dtype={OOCYTE: str})
    else:
        print("NOT FOUND!!!")
        print(csv_file_path)
    return None


def compare_masks(reference: str, alternatives: str, measures=all_measures,
                  in_csv_metrics="eval.csv", out_csv_metrics="out.csv"):
    df_old = load_csv(in_csv_metrics)
    print(df_old)
    reference_dataset = aix.Dataset.from_file(AIX_DATASETS / reference)
    alt_datasets = {}
    for alt in alternatives:
        alt_datasets[alt] = aix.Dataset.from_file(AIX_DATASETS / alt)
    l = []
    for item_mask in tqdm(reference_dataset.iterate_mask_items()):
        ref_mask_img = item_mask.norm_image(cv2.IMREAD_GRAYSCALE)
        for alt, alt_d in alt_datasets.items():
            if item_mask.index in alt_d.oocytes:
                alt_mask = aix.Item(alt_d, mask=item_mask.mask, stage=item_mask.stage, index=item_mask.index)
                alt_mask_img = alt_mask.norm_image(cv2.IMREAD_GRAYSCALE)
                d = {OOCYTE: item_mask.index,
                     MATURE: item_mask.stage,
                     REFERENCE: reference_dataset.name,
                     COMPARED: alt_d.name}
                for measure_name, measure in measures.items():
                    m = measure(ref_mask_img, alt_mask_img)
                    d[measure_name] = m
                l.append(d)
            else:
                print("Ignoring oocyte " + str(item_mask.index) + " for dataset " + alt_d.name)
        #break
    df = pd.DataFrame(l)
    if df_old is not None:
        index_fields = [OOCYTE, MATURE, REFERENCE, COMPARED]
        df = df_old.merge(df, how='outer', on=index_fields, suffixes=("_x", None))
        print("After merge")
        print(df[COMPARED].unique())
        for measure_name in measures:
            index_fields.append(measure_name)
            df[measure_name].fillna(df[measure_name + "_x"], inplace=True)
        df = df[index_fields]
    out_csv_file_path = AIX_EVALS / out_csv_metrics
    out_csv_file_path.parent.mkdir(parents=True, exist_ok=True)
    print(out_csv_file_path)
    df.to_csv(out_csv_file_path, index=False)

    return df


def make_dict(df: pd.DataFrame):
    d = {}

    metrics = df.columns
    for index, row in df.iterrows():
        name = "-".join(index)
        for i, val in enumerate(row):
            namem = name + "-" + metrics[i]
            d[namem] = val
    return d

def summarize_comparison(df:pd.DataFrame, metrics_file="eval.json"):
    print(df)
    gb = df.drop([MATURE,OOCYTE], axis=1).groupby([REFERENCE, COMPARED])
    print(gb)
    means = make_dict(gb.mean())
    medians = make_dict(gb.median())
    d = {"MEANS": means, "MEDIANS": medians}
    filepath = AIX_EVALS/metrics_file
    with filepath.open("w") as f:
        json.dump(d, f)
    # df.groupby([REFERENCE, COMPARED, MATURE])
    # df.mean()
    # df.median()
    #
    # measures = df.columns[4:]
    # labelings = np.unique(df[])
    # for alt in alternatives:
    #     print("==", alt)
    #
    #     for measure_name in measures:
    #         condition = (df["reference"] == reference_dataset.name) & (df["compared"] == alt_datasets[alt].name)
    #         measure_slice = df[condition][measure_name].to_numpy()
    #         print("   Mean " + measure_name + ": %4.4f" % np.mean(measure_slice))
    #         print("   Median " + measure_name + ": %4.4f" % np.median(measure_slice))
    #         for mature in [False,True]:
    #             condition_mature = condition & (df[MATURE]==mature)
    #             measure_slice = df[condition_mature][measure_name].to_numpy()
    #             mature_txt = MATURE if mature else IMMATURE
    #             print("   Mean " + measure_name + " " + mature_txt + ": %4.4f" % np.mean(measure_slice))
    #             print("   Median " + measure_name + " " + mature_txt + ": %4.4f" % np.median(measure_slice))

            
            
def evaluate_scale_loss(reference: str, dim, interpolation=cv2.INTER_AREA, extrapolation=cv2.INTER_CUBIC, measures=all_measures):
    reference_dataset = aix.Dataset.from_file(reference)
    measure_vals = {}
    for measure_name in measures:
        measure_vals[measure_name] = []
    for item_mask in tqdm(reference_dataset.iterate_mask_items()):
        ref_mask_img = item_mask.uint_norm_image()
        
        resized = cv2.resize(ref_mask_img, dim, interpolation=interpolation)
        back_area = cv2.resize(resized, (ref_mask_img.shape[1], ref_mask_img.shape[0]), interpolation=extrapolation)
        back_area[back_area<0.]=0.
        back_area[back_area>1.]=1.
        for measure_name, measure in measures.items():
            measure_vals[measure_name].append(measure(ref_mask_img, back_area))
         
    print("\nReference: " + reference + "\n")
    
    first_measure_name = next(iter(measures))
    n_masks = len(measure_vals[first_measure_name])
    print("   N.Images : %4d" % n_masks)
    for measure_name in measures:
        print("   " + measure_name + ": %4.2f" % (sum(measure_vals[measure_name]) / n_masks))