import numpy as np
import pandas as pd
from collections import Counter
from sklearn.linear_model import SGDClassifier
from sklearn.metrics import accuracy_score, f1_score, precision_score, recall_score, classification_report, confusion_matrix
from sklearn.preprocessing import StandardScaler, PolynomialFeatures
import logging
import os
from datetime import datetime
import argparse
import joblib
import json
from sklearn.impute import SimpleImputer, KNNImputer
import zarr
import wandb
import time
import psutil


# Parse command-line arguments
parser = argparse.ArgumentParser(description="Train a model using land use data.")
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("--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")
parser.add_argument("--batch_size", type=int, default=500, help="Batch size for incremental learning")

args = parser.parse_args()
start_year = args.start_year
end_year = args.end_year
region = args.region
imputation_method = args.imputation
train_years = args.train_years
test_years = args.test_years
batch_size = args.batch_size

# Initialize W&B
wandb.init(
    project="land-use-classification",
    config={
        "start_year": start_year,
        "end_year": end_year,
        "region": region,
        "imputation": imputation_method,
        "train_years": train_years,
        "test_years": test_years,
        "batch_size": batch_size,
    },
    mode="offline"
)

# Function to convert numpy data types to standard Python types
def convert_numpy_types(obj):
    if isinstance(obj, np.integer):
        return int(obj)
    elif isinstance(obj, np.floating):
        return float(obj)
    elif isinstance(obj, np.ndarray):
        return obj.tolist()  # Convert numpy arrays to lists
    else:
        return obj

# 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("Training with SGDClassifier")

# Set paths for saving/loading processed data and metadata
save_data_path = '/gpfs/projects/bsc32/cerise/outputs/bsc032013/ml/'
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)
logging.info(f"Loaded metadata: {metadata}")

def load_data_for_years(years, region, data_path):
    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.")

# Initialize scaler
scaler = StandardScaler()

# Initialize the SGDClassifier for incremental learning
sgd_clf = SGDClassifier(loss="log_loss", max_iter=1000, tol=1e-3)

# Scaling training data
X_train = scaler.fit_transform(X_train)
X_test = scaler.transform(X_test)

# Initialize PolynomialFeatures with degree 2 to add interaction terms
poly = PolynomialFeatures(degree=2)

# Transform the features to add nonlinearity
X_train_poly = poly.fit_transform(X_train)
X_test_poly = poly.transform(X_test)

# Train the model on transformed data
sgd_clf.partial_fit(X_train_poly, y_train, classes=np.unique(y_train))

# Train in batches
num_batches = X_train.shape[0] // batch_size + 1

# 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)}

# 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])

# Batch training with incremental learning
for batch in range(num_batches):
    start_idx = batch * batch_size
    end_idx = min(start_idx + batch_size, X_train.shape[0])
    
    X_batch = X_train[start_idx:end_idx]
    y_batch = y_train[start_idx:end_idx]

    sgd_clf.partial_fit(X_batch, y_batch, classes=np.unique(y_train))

    logging.info(f"Trained batch {batch+1}/{num_batches}")

# Evaluate model on test data
y_pred = sgd_clf.predict(X_test)

# Compute performance metrics
accuracy = accuracy_score(y_test, y_pred)
f1 = f1_score(y_test, y_pred, average='weighted')
precision = precision_score(y_test, y_pred, average='weighted')
recall = recall_score(y_test, y_pred, average='weighted')
conf_matrix = confusion_matrix(y_test, y_pred)

# Log the performance metrics to W&B
wandb.log({
    'accuracy': accuracy,
    'f1_score': f1,
    'precision': precision,
    'recall': recall,
    'confusion_matrix': conf_matrix.tolist()
})

# Log performance
logging.info(f'Accuracy: {accuracy*100:.2f}%')
logging.info(f'F1 Score: {f1:.2f}')
logging.info(f'Precision: {precision:.2f}')
logging.info(f'Recall: {recall:.2f}')
logging.info(f'Confusion Matrix:\n{conf_matrix}')

# Log detailed classification report
class_report = classification_report(y_test, y_pred, target_names=[f'Class {i}' for i in range(len(unique_classes))])
logging.info(f'Classification Report:\n{class_report}')

# Save the trained model
model_filename = f'SGDClassifier_{start_year}_{end_year}_{region}_{timestamp}.joblib'
joblib.dump(sgd_clf, os.path.join(save_data_path, model_filename))
logging.info(f'SGDClassifier model saved successfully to {model_filename}')

# Save metadata
training_metadata = {
    'start_year': start_year,
    'end_year': end_year,
    'region': region,
    'batch_size': batch_size,
    'class_to_index': {str(k): v for k, v in class_to_index.items()},
    'X_train_shape': X_train.shape,
    'X_test_shape': X_test.shape,
    'train_class_distribution': dict(Counter(y_train)),
    'test_class_distribution': dict(Counter(y_test)),
}

training_metadata = {k: convert_numpy_types(v) for k, v in training_metadata.items()}

# Save metadata
metadata_filename = f'SGDClassifier_metadata_{start_year}_{end_year}_{region}_{timestamp}.json'
metadata_path = os.path.join(save_data_path, metadata_filename)

with open(metadata_path, 'w') as f:
    json.dump(training_metadata, f, indent=4)
logging.info(f'Training metadata saved successfully to {metadata_path}')
