import numpy as np
import pandas as pd
from collections import Counter
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 argparse
import joblib
import json
from sklearn.impute import SimpleImputer, KNNImputer  
from sklearn.utils.class_weight import compute_class_weight
import zarr
import sqlite3
import json
import sys 

sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../source')))
from database import log_model_run
from data_processing import apply_land_sea_mask
from train import compute_weights, train_evaluate_ml_model

# Parse command-line arguments
parser = argparse.ArgumentParser(description="Train a model using land use data.")
parser.add_argument("--data_path", type=str, required=True, help="Path to the directory where data is stored")
parser.add_argument("--db_path", type=str, required=True, help="Path to the SQLite database for logging model runs")
parser.add_argument("--start_year", type=int, required=True, help="Start year for the data range")
parser.add_argument("--end_year", type=int, required=True, help="End year for the data range")
parser.add_argument("--region", type=str, required=True, help="Region for the data slice (ES, EU, UK, AU, GLOBE)")
parser.add_argument("--model_type", type=str, required=True, choices=['xgboost', 'logistic_regression'], help="Model type to use")
parser.add_argument("--imputation", type=str, choices=['mean', 'median', 'most_frequent', 'knn'], default='mean', help="Imputation method for handling NaN values")
parser.add_argument("--train_years", nargs='+', type=int, required=True, help="Years to use for training")
parser.add_argument("--test_years", nargs='+', type=int, required=True, help="Years to use for testing")

args = parser.parse_args()
start_year = args.start_year
end_year = args.end_year
region = args.region
model_type = args.model_type
imputation_method = args.imputation
train_years = args.train_years
test_years = args.test_years
save_data_path = args.data_path
db_path= args.db_path

# Configure logging
current_path = os.getcwd()
logs_path = os.path.join(current_path, "Logs")
os.makedirs(logs_path, exist_ok=True)
timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
log_file_name = f"train_log_{start_year}_{end_year}_{region}_{timestamp}.log"
log_file_path = os.path.join(logs_path, log_file_name)
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s', handlers=[
    logging.FileHandler(log_file_path),
    logging.StreamHandler()
])
logging.info(f"Training model of type: {model_type}")

# Set paths for saving/loading processed data and metadata
metadata_path = os.path.join(save_data_path, f'metadata_{start_year}_{end_year}_{region}.json')

# Load metadata
with open(metadata_path, 'r') as f:
    metadata = json.load(f)

# Log the loaded metadata
logging.info(f"Loaded metadata: {metadata}")

def load_data_for_years(years, region, data_path):
    """
    Load features and lccs_class data for the specified years.

    :param years: list of int, Years to load data for
    :param region: str, Region code (ES, EU, UK, AU, GLOBE)
    :param data_path: str, Base path where data is stored
    :return: tuple of np.array, Concatenated features and lccs_class for the specified years
    """
    features_list = []
    lccs_class_list = []

    for year in years:
        features_year_path = os.path.join(data_path, f'features_{year}_{region}.zarr')
        lccs_class_year_path = os.path.join(data_path, f'lccs_class_{year}_{region}.zarr')

        features_list.append(zarr.load(features_year_path))
        lccs_class_list.append(zarr.load(lccs_class_year_path))

    features = np.concatenate(features_list, axis=0)
    lccs_class = np.concatenate(lccs_class_list, axis=0)

    return features, lccs_class

# Load data for training and testing years
X_train, y_train = load_data_for_years(train_years, region, save_data_path)
X_test, y_test = load_data_for_years(test_years, region, save_data_path)

# Impute missing values
if imputation_method in ['mean', 'median', 'most_frequent']:
    imputer = SimpleImputer(strategy=imputation_method)
else:
    imputer = KNNImputer(n_neighbors=5)

X_train = imputer.fit_transform(X_train)
X_test = imputer.transform(X_test)
logging.info(f"Imputed missing values using {imputation_method} strategy.")

# Adjust labels to ensure they start from 0 and go up to n_classes-1
unique_classes = sorted(np.unique(y_train))
class_to_index = {float(cls): idx for idx, cls in enumerate(unique_classes)}  # Convert keys to float

# Save class_to_index to a JSON file
with open(os.path.join(save_data_path, f'class_to_index_{start_year}_{end_year}_{region}_{timestamp}.json'), 'w') as f:
    json.dump(class_to_index, f)

logging.info("Class-to-index mapping saved successfully.")

y_train = np.array([class_to_index[cls] for cls in y_train])
y_test = np.array([class_to_index[cls] for cls in y_test])

# Check Data Balance
class_counts = Counter(y_train)
class_counts_df = pd.DataFrame.from_dict(class_counts, orient='index', columns=['count'])
logging.info("Class distribution in training set:")
logging.info(class_counts_df.sort_index())


class_weight_dict, sample_weight = compute_weights(y_train)

print("Class Weight Dictionary:", class_weight_dict)
print("Sample Weights:", sample_weight)


# Initialize and train the model based on the chosen architecture
if model_type == "xgboost":
    model = xgb.XGBClassifier(
        n_estimators=2000,           
        max_depth=10,                
        learning_rate=0.04,          
        subsample=0.8,               
        reg_alpha=0.5,               
        reg_lambda=1.0,              
        colsample_bytree=0.8,        
        eval_metric='mlogloss',      
        tree_method='hist',         
        device='cuda',              
        random_state=42             
    )

    # Capture returned metrics
    #accuracy, f1, precision, recall, feature_importances = train_evaluate_ml_model(model, X_train, X_test, y_train, y_test, class_to_index, "XGBoost")
    accuracy, f1, precision, recall, feature_importances = train_evaluate_ml_model(
            model=model,
            X_train=X_train, 
            X_test=X_test, 
            y_train=y_train, 
            y_test=y_test, 
            class_to_index=class_to_index, 
            model_name='XGBoost', 
            sample_weight=sample_weight, 
            save_data_path=save_data_path, 
            start_year=start_year, 
            end_year=end_year, 
            region=region, 
            timestamp=timestamp
        )
    
elif model_type == "logistic_regression":
    model = LogisticRegression(max_iter=2000, solver='saga', random_state=42, class_weight='balanced')

    # Capture returned metrics
    accuracy, f1, precision, recall, feature_importances = train_evaluate_ml_model(model, X_train, X_test, y_train, y_test, class_to_index, "LogisticRegression")

else:
    raise ValueError("Invalid model type specified. Choose from 'xgboost', 'logistic_regression'.")

# Now, you can use the returned metrics and feature_importances in log_model_run
model_file_path = os.path.join(save_data_path, f'{model_type}_{start_year}_{end_year}_{region}_{timestamp}.joblib')

log_model_run(db_path=args.db_path, 
              model_type=model_type, 
              data_path=save_data_path, 
              region=region, 
              start_year=start_year, 
              end_year=end_year, 
              train_years=train_years, 
              test_years=test_years, 
              timestamp=timestamp, 
              accuracy=accuracy, 
              f1=f1, 
              precision=precision, 
              recall=recall, 
              feature_importance=feature_importances,  # No need for conditional check anymore
              class_mapping=class_to_index, 
              model_file_path=model_file_path)


logging.info("Training and evaluation complete.")
