import numpy as np
import pandas as pd
import xgboost as xgb  
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score, f1_score, precision_score, recall_score, classification_report
import logging
import os
from datetime import datetime
import joblib
from sklearn.impute import SimpleImputer, KNNImputer  
from sklearn.utils.class_weight import compute_class_weight
import json
import sys

sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../source')))
from data_processing import apply_land_sea_mask

def compute_weights(y_train):
    """
    Computes class weights and sample weights for balancing the dataset.
    
    Parameters:
    y_train (array-like): The training labels.

    Returns:
    class_weight_dict (dict): A dictionary mapping each class to its computed weight.
    sample_weight (np.array): An array of sample weights corresponding to the training labels.
    """
    # Compute class weights
    class_weight = compute_class_weight(class_weight='balanced', classes=np.unique(y_train), y=y_train)
    class_weight_dict = {cls: weight for cls, weight in zip(np.unique(y_train), class_weight)}
    
    # Compute sample weights
    sample_weight = np.array([class_weight_dict[cls] for cls in y_train])
    
    logging.info("Computed class weights for balancing.")
    
    return class_weight_dict, sample_weight


# Function to train and evaluate ML models with performance metrics, feature importance, and model saving
def train_evaluate_ml_model(model, X_train, X_test, y_train, y_test, class_to_index, model_name, 
                            sample_weight, save_data_path, start_year, end_year, region, timestamp):
    # Fit the model
    model.fit(X_train, y_train, sample_weight=sample_weight)
    y_pred = model.predict(X_test)

    # Apply the water-land mask to predictions (if applicable)
    y_pred_adjusted = apply_land_sea_mask(y_pred, X_test, class_to_index)

    # Compute performance metrics
    accuracy = accuracy_score(y_test, y_pred_adjusted)
    f1 = f1_score(y_test, y_pred_adjusted, average='weighted')
    precision = precision_score(y_test, y_pred_adjusted, average='weighted', zero_division=0)
    recall = recall_score(y_test, y_pred_adjusted, average='weighted')

    # Log the performance metrics
    logging.info(f'Accuracy of {model_name}: {accuracy*100:.2f}%')
    logging.info(f'F1 Score of {model_name}: {f1:.2f}')
    logging.info(f'Precision of {model_name}: {precision:.2f}')
    logging.info(f'Recall of {model_name}: {recall:.2f}')

    # Generate the correct number of target names
    unique_classes_in_test = np.unique(np.concatenate([y_test, y_pred_adjusted]))
    target_names = [f'Class {i}' for i in unique_classes_in_test]

    # Log detailed classification report
    class_report = classification_report(y_test, y_pred_adjusted, target_names=target_names, zero_division=0)
    logging.info(f'Classification Report for {model_name}:\n{class_report}')

    # Log feature importance for XGBoost
    feature_importances = None  # Initialize in case it's not available
    if hasattr(model, "feature_importances_"):
        feature_importances = model.feature_importances_
        importance_df = pd.DataFrame({
            'Feature': np.arange(len(feature_importances)),
            'Importance': feature_importances
        }).sort_values(by='Importance', ascending=False)

        logging.info(f'Feature Importances for {model_name}:\n{importance_df}')

    # Log coefficients for Logistic Regression
    elif model_name == "LogisticRegression" and hasattr(model, 'coef_'):
        feature_importances = model.coef_[0]  # Assuming binary classification; adapt for multiclass
        importance_df = pd.DataFrame({
            'Feature': np.arange(len(feature_importances)),
            'Importance': feature_importances
        }).sort_values(by='Importance', ascending=False)

        logging.info(f'Feature Coefficients for {model_name}:\n{importance_df}')

    # Save the model
    model_filename = f'{model_name}_{start_year}_{end_year}_{region}_{timestamp}.joblib'
    model_path = os.path.join(save_data_path, model_filename)
    joblib.dump(model, model_path)
    logging.info(f'{model_name} model saved successfully to {model_path}.')
    
    return accuracy, f1, precision, recall, feature_importances
