import torch
import torch.nn as nn
import numpy as np
import pandas as pd
import pandas.api.types
import sklearn.metrics
import warnings

# Suppress the specific FutureWarning
warnings.simplefilter(action='ignore', category=FutureWarning)

# Your custom metric function
def metric(y_true, y_pred):
    d_prefix = ["_left_neural_foraminal_narrowing_l1_l2",
                "_left_neural_foraminal_narrowing_l2_l3",
                "_left_neural_foraminal_narrowing_l3_l4",
                "_left_neural_foraminal_narrowing_l4_l5",
                "_left_neural_foraminal_narrowing_l5_s1",
                "_right_neural_foraminal_narrowing_l1_l2",
                "_right_neural_foraminal_narrowing_l2_l3",
                "_right_neural_foraminal_narrowing_l3_l4",
                "_right_neural_foraminal_narrowing_l4_l5",
                "_right_neural_foraminal_narrowing_l5_s1",
                '_left_subarticular_stenosis_l1_l2',
                '_left_subarticular_stenosis_l2_l3',
                '_left_subarticular_stenosis_l3_l4',
                '_left_subarticular_stenosis_l4_l5',
                '_left_subarticular_stenosis_l5_s1',
                '_right_subarticular_stenosis_l1_l2',
                '_right_subarticular_stenosis_l2_l3',
                '_right_subarticular_stenosis_l3_l4',
                '_right_subarticular_stenosis_l4_l5',
                '_right_subarticular_stenosis_l5_s1',
                "_spinal_canal_stenosis_l1_l2",
                "_spinal_canal_stenosis_l2_l3",
                "_spinal_canal_stenosis_l3_l4",
                "_spinal_canal_stenosis_l4_l5",
                "_spinal_canal_stenosis_l5_s1"]
    solution = pd.DataFrame({'row_id': [], 'normal_mild': [], 'moderate': [], 'severe': [], 'sample_weight': []})
    submission = pd.DataFrame({'row_id': [], 'normal_mild': [], 'moderate': [], 'severe': []})
    y_true = np.eye(3)[y_true.detach().numpy()]
    y_pred = y_pred.detach().numpy()
    for k, i in enumerate(range(y_true.shape[0])):
        d_tmp = [str(i) + l for l in d_prefix]
        solution_tmp = pd.DataFrame({'row_id': d_tmp, 'normal_mild': y_true[0, :, 0], 'moderate': y_true[0, :, 1], 'severe': y_true[0, :, 2]})
        submission_tmp = pd.DataFrame({'row_id': d_tmp, 'normal_mild': y_pred[0, :, 0], 'moderate': y_pred[0, :, 1], 'severe': y_pred[0, :, 2]})
        solution_tmp['sample_weight'] = 0
        solution_tmp['sample_weight'] = solution_tmp.apply(lambda x: 1 if x['normal_mild'] == 1 else (2 if x['moderate'] == 1 else 4), axis=1)

        solution = pd.concat([solution, solution_tmp])
        submission = pd.concat([submission, submission_tmp])
    return score(solution, submission, 'row_id', 1.0)

class ParticipantVisibleError(Exception):
    pass

def get_condition(full_location: str) -> str:
    # Given an input like spinal_canal_stenosis_l1_l2 extracts 'spinal'
    for injury_condition in ['spinal', 'foraminal', 'subarticular']:
        if injury_condition in full_location:
            return injury_condition
    raise ValueError(f'condition not found in {full_location}')

def score(solution: pd.DataFrame, submission: pd.DataFrame, row_id_column_name: str, any_severe_scalar: float) -> float:
    target_levels = ['normal_mild', 'moderate', 'severe']


    solution['study_id'] = solution['row_id'].apply(lambda x: x.split('_')[0])
    solution['location'] = solution['row_id'].apply(lambda x: '_'.join(x.split('_')[1:]))
    solution['condition'] = solution['row_id'].apply(get_condition)

    del solution[row_id_column_name]
    del submission[row_id_column_name]
    assert sorted(submission.columns) == sorted(target_levels)

    submission['study_id'] = solution['study_id']
    submission['location'] = solution['location']
    submission['condition'] = solution['condition']

    condition_losses = []
    condition_weights = []
    for condition in ['spinal', 'foraminal', 'subarticular']:
        condition_indices = solution.loc[solution['condition'] == condition].index.values
        condition_loss = sklearn.metrics.log_loss(
            y_true=solution.loc[condition_indices, target_levels].values,
            y_pred=submission.loc[condition_indices, target_levels].values,
            sample_weight=solution.loc[condition_indices, 'sample_weight'].values
        )
        condition_losses.append(condition_loss)
        condition_weights.append(1 / solution.loc[condition_indices, 'location'].nunique())

    any_severe_spinal_labels = pd.Series(solution.loc[solution['condition'] == 'spinal'].groupby('study_id')['severe'].max())
    any_severe_spinal_weights = pd.Series(solution.loc[solution['condition'] == 'spinal'].groupby('study_id')['sample_weight'].max())
    
    any_severe_spinal_predictions = pd.Series(submission.loc[submission['condition'] == 'spinal'].groupby('study_id')['severe'].max())
    any_severe_spinal_loss = sklearn.metrics.log_loss(
        y_true=any_severe_spinal_labels,
        y_pred=any_severe_spinal_predictions,
        sample_weight=any_severe_spinal_weights, labels=[0, 1]
    )
    condition_losses.append(any_severe_spinal_loss)
    condition_weights.append(any_severe_scalar)
    return np.average(condition_losses, weights=condition_weights)