import sys
sys.path.append('.')
from tqdm import tqdm
import torch
from torch.nn import functional as F
from .metrics import roc_auc, pr_auc, calc_metrics
import os
import yaml
import numpy as np


def evaluate(conf):
    device = conf['device']
    dataloader = conf['dataloaders']['test']
    experiment_dir = conf['experiment_dir']
    classes = conf['data']['classes']
    batch_size = conf['data']['batch_size']

    model = conf['model']
    model.load_state_dict(conf['best_weights'])
    model = model.to(device)
    model.eval()


    ground_truth = None
    inferences = None

    batch_bar = tqdm(dataloader, desc='Batch', unit='batches', leave=True)
    for inputs, labels in batch_bar:
        inputs = inputs.to(device)
        with torch.set_grad_enabled(False):
            outputs = model(inputs)

        probs = F.softmax(outputs, dim=1)[:, 1]
        probs = probs.cpu()
        if ground_truth is None and inferences is None:
            ground_truth = labels
            inferences = probs
        else:
            ground_truth = torch.cat((ground_truth, labels))
            inferences = torch.cat((inferences, probs))

        # Calculate save metrics

    metrics = {'metrics0.5': calc_metrics(ground_truth=ground_truth,
                                          inferences=inferences,
                                          threshold=0.5),
               'metrics0.7': calc_metrics(ground_truth=ground_truth,
                                          inferences=inferences,
                                          threshold=0.7),
               'metrics0.9': calc_metrics(ground_truth=ground_truth,
                                          inferences=inferences,
                                          threshold=0.9),
               'roc_auc': roc_auc(ground_truth=ground_truth,
                                  inferences=inferences,
                                  experiment_dir=experiment_dir),
               'pr_auc': pr_auc(ground_truth=ground_truth,
                                inferences=inferences,
                                experiment_dir=experiment_dir),
                }

    patients_dataset = conf['patients_dataset']
    test_patients = patients_dataset.test_patients
    patients_bar = tqdm(test_patients.items(),
                        desc='Patient', unit='patients', leave=True)

    inferences = {1: [], 2:[], 3: [], 4: [], 5: [],
                 '5%':[], '7.5%': [], '10%': []}
    ground_truth = []
    for patient, patient_data in patients_bar:
        patient_IH = patient_data['IH'] # If the patient has IH or not
        slices = patient_data['slices_IH'] + patient_data['slices_noIH']
        slices_with_IH = 0
        samples = []
        for slice_id in slices:
            sample, _ = patients_dataset.getSlice(slice_id)
            samples.append(sample)

        for i in range(0, len(samples), batch_size):
            batch = torch.stack(samples[i:i+batch_size], dim=0)
            batch = batch.to(device)
            with torch.set_grad_enabled(False):
                outputs = model(batch)
            IH_probs = F.softmax(outputs, dim=1)[:, 1]
            slices_with_IH += (IH_probs > 0.8).sum()
        for num_IH_threshold in [1,2,3,4,5]:
            net_IH_prediction = slices_with_IH >= num_IH_threshold
            inferences[num_IH_threshold].append(net_IH_prediction)
        for percentage, key in [(0.05, '5%'), (0.075, '7.5%'), (0.10, '10%')]:
            num_IH_threshold = max(1, round(percentage * len(slices)))
            net_IH_prediction = slices_with_IH >= num_IH_threshold
            inferences[key].append(net_IH_prediction)

        ground_truth.append(patient_IH)


    ground_truth = np.array(ground_truth).astype(float)
    inferences1 = np.array(inferences[1]).astype(float)
    inferences2 = np.array(inferences[2]).astype(float)
    inferences3 = np.array(inferences[3]).astype(float)
    inferences4 = np.array(inferences[4]).astype(float)
    inferences5 = np.array(inferences[5]).astype(float)
    inferencesperc1 = np.array(inferences['5%']).astype(float)
    inferencesperc2 = np.array(inferences['7.5%']).astype(float)
    inferencesperc3 = np.array(inferences['10%']).astype(float)
    metrics['patients_metrics (>= 1 IH slice)'] = calc_metrics(
                                                    ground_truth=ground_truth,
                                                    inferences=inferences1)
    metrics['patients_metrics (>= 2 IH slice)'] = calc_metrics(
                                                    ground_truth=ground_truth,
                                                    inferences=inferences2)
    metrics['patients_metrics (>= 3 IH slice)'] = calc_metrics(
                                                    ground_truth=ground_truth,
                                                    inferences=inferences3)
    metrics['patients_metrics (>= 4 IH slice)'] = calc_metrics(
                                                    ground_truth=ground_truth,
                                                    inferences=inferences4)
    metrics['patients_metrics (>= 5 IH slice)'] = calc_metrics(
                                                    ground_truth=ground_truth,
                                                    inferences=inferences5)
    metrics['patients_metrics (>= 5% IH slice)'] = calc_metrics(
                                                    ground_truth=ground_truth,
                                                    inferences=inferencesperc1)
    metrics['patients_metrics (>= 7.5% IH slice)'] = calc_metrics(
                                                    ground_truth=ground_truth,
                                                    inferences=inferencesperc2)
    metrics['patients_metrics (>= 10% IH slice)'] = calc_metrics(
                                                    ground_truth=ground_truth,
                                                    inferences=inferencesperc3)

    del metrics['patients_metrics (>= 1 IH slice)']['threshold']
    del metrics['patients_metrics (>= 2 IH slice)']['threshold']
    del metrics['patients_metrics (>= 3 IH slice)']['threshold']
    del metrics['patients_metrics (>= 4 IH slice)']['threshold']
    del metrics['patients_metrics (>= 5 IH slice)']['threshold']
    del metrics['patients_metrics (>= 5% IH slice)']['threshold']
    del metrics['patients_metrics (>= 7.5% IH slice)']['threshold']
    del metrics['patients_metrics (>= 10% IH slice)']['threshold']

    with open(os.path.join(experiment_dir, 'metrics.yaml'), 'w') as fp:
        yaml.dump(metrics, fp)