import pandas as pd
import numpy as np
from sklearn.preprocessing import OneHotEncoder, StandardScaler
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset
from torch import optim
from sklearn.model_selection import train_test_split

def data_loader(train_path, test_path):
    # Load training and test data
    train_data = pd.read_excel(train_path)
    test_data = pd.read_excel(test_path)

    # Identify features and target variable
    features = ['x', 'Annual Mean Period', 'Annual Mean Spring Tidal Range', 
                'Annual Mean Tidal Range', 'Breaker Wave Height Hb', 
                'Deep Water Wave Height Hd', 'Mean Wave Height', 
                'Dimensionless Settling Velocity', 
                'High Tide Sediment Settling Velocity', 
                'Latitude', 'Longitude', 'Mean Grain Size', 
                'Mean Grain Size (Mz)', 'Intertidal Slope', 
                'Skewness', 'Kurtosis', 'Sorting Coefficient']
    target = 'y'
    categorical = ['Dominant Wave Direction']

    # One-hot encode the categorical variable
    encoder = OneHotEncoder(sparse=False)
    train_encoded = encoder.fit_transform(train_data[categorical])
    test_encoded = encoder.transform(test_data[categorical])

    # Normalize the numerical features
    scaler = StandardScaler()
    train_scaled = scaler.fit_transform(train_data[features])
    test_scaled = scaler.transform(test_data[features])

    # Combine encoded categorical and scaled numerical features
    train_features = np.hstack([train_scaled, train_encoded])
    test_features = np.hstack([test_scaled, test_encoded])

    # Prepare the final DataFrames
    train_df = pd.DataFrame(train_features, columns=features + list(encoder.get_feature_names(categorical)))
    train_df[target] = train_data[target]

    test_df = pd.DataFrame(test_features, columns=features + list(encoder.get_feature_names(categorical)))
    test_df[target] = test_data[target]

    return train_df, test_df

def theoretical_model_integration(train_df, test_df):
    # Example coefficients for derived feature computation from Bruun and Dean models
    S = 0.1  # Subject to real-world data
    d = 5  # Depth of closure
    B = 2  # Berm height
    a = 8  # Active height
    A = 0.5  # Dean model coefficient
    
    # Calculate derived feature from Bruun model
    train_df['retreat_distance'] = (S * (d + B)) / (a + B)
    test_df['retreat_distance'] = (S * (d + B)) / (a + B)
    
    # Calculate derived feature from Dean model
    train_df['dean_profile'] = A * (train_df['x'] ** (2/3))
    test_df['dean_profile'] = A * (test_df['x'] ** (2/3))

    return train_df, test_df

def deep_learning_model(X_train, X_test):
    # Define neural network architecture
    class NeuralNetwork(nn.Module):
        def __init__(self, input_dim):
            super(NeuralNetwork, self).__init__()
            self.layer1 = nn.Linear(input_dim, 64)
            self.layer2 = nn.Linear(64, 64)
            self.layer3 = nn.Linear(64, 1)
        
        def forward(self, x):
            x = torch.relu(self.layer1(x))
            x = torch.relu(self.layer2(x))
            x = self.layer3(x)
            return x

    # Initialize model
    input_dim = X_train.shape[1]
    model = NeuralNetwork(input_dim)
    
    # Compile the model
    loss_fn = nn.MSELoss()
    optimizer = optim.Adam(model.parameters(), lr=0.001)

    return model, loss_fn, optimizer

def model_training_evaluation(model, loss_fn, optimizer, train_df, test_df):
    # Extract features and target
    X_train = train_df.drop(columns=['y']).values
    y_train = train_df['y'].values
    X_test = test_df.drop(columns=['y']).values
    y_test = test_df['y'].values

    # Split train data into train and validation
    X_train_split, X_val, y_train_split, y_val = train_test_split(X_train, y_train, test_size=0.2, random_state=42)

    # Convert data into Tensor
    train_dataset = TensorDataset(torch.tensor(X_train_split, dtype=torch.float32), torch.tensor(y_train_split, dtype=torch.float32))
    val_dataset = TensorDataset(torch.tensor(X_val, dtype=torch.float32), torch.tensor(y_val, dtype=torch.float32))
    test_dataset = TensorDataset(torch.tensor(X_test, dtype=torch.float32), torch.tensor(y_test, dtype=torch.float32))

    # DataLoader
    train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
    val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False)
    test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)

    # Train the model
    epochs = 100
    for epoch in range(epochs):
        model.train()
        for X_batch, y_batch in train_loader:
            optimizer.zero_grad()
            predictions = model(X_batch).squeeze()
            loss = loss_fn(predictions, y_batch)
            loss.backward()
            optimizer.step()

        # Validation step
        model.eval()
        val_loss_total = 0
        with torch.no_grad():
            for X_val_batch, y_val_batch in val_loader:
                val_predictions = model(X_val_batch).squeeze()
                val_loss = loss_fn(val_predictions, y_val_batch)
                val_loss_total += val_loss.item()

    # Evaluate on test set
    model.eval()
    test_loss_total = 0
    y_preds = []
    with torch.no_grad():
        for X_test_batch, y_test_batch in test_loader:
            test_predictions = model(X_test_batch).squeeze()
            test_loss = loss_fn(test_predictions, y_test_batch)
            test_loss_total += test_loss.item()
            y_preds.extend(test_predictions.numpy())

    # Calculate additional evaluation metrics
    rmse = np.sqrt(test_loss_total / len(test_loader))
    mae = np.mean(np.abs(np.array(y_preds) - y_test))
    r_squared = 1 - (np.sum((np.array(y_preds) - y_test) ** 2) / np.sum((y_test - y_test.mean()) ** 2))

    return {'RMSE': rmse, 'MAE': mae, 'R^2': r_squared}

def main_function(train_path, test_path):
    # Load and preprocess data
    train_df, test_df = data_loader(train_path, test_path)

    # Integrate theoretical model features
    train_df, test_df = theoretical_model_integration(train_df, test_df)

    # Prepare input and output data for model
    X_train = train_df.drop(columns=['y']).values
    X_test = test_df.drop(columns=['y']).values

    # Define deep learning model
    model, loss_fn, optimizer = deep_learning_model(X_train, X_test)

    # Train and evaluate the model
    evaluation_metrics = model_training_evaluation(model, loss_fn, optimizer, train_df, test_df)

    # Return evaluation results
    return evaluation_metrics