import torch
import numpy as np
import matplotlib.pyplot as plt
from Dataset import Dataset

# DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
DEVICE = torch.device('cpu')

# Set global plotting parameters
plt.rcParams.update({'font.size': 14,
                     'figure.figsize': (10, 8),
                     'lines.linewidth':  2,
                     'lines.markersize': 6,
                     'axes.grid': True,
                     'axes.labelsize': 16,
                     'legend.fontsize': 14,
                     'xtick.labelsize': 14,
                     'ytick.labelsize': 14,
                     'figure.autolayout': True
                     })

def set_seed(seed=42):
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)    

class NeuralNetwork(torch.nn.Module):
    def __init__(self, layer_sizes, dropout_rate=0.0, activation=torch.nn.ReLU):
        super(NeuralNetwork, self).__init__()

        if dropout_rate > 0:
            self.dropout_layer = torch.nn.Dropout(dropout_rate)

        self.layer_sizes = layer_sizes
        self.layers = torch.nn.ModuleList()
        for i in range(len(layer_sizes) - 2):
            self.layers.append(torch.nn.Linear(layer_sizes[i], layer_sizes[i + 1]))
            self.layers.append(activation())
        self.layers.append(torch.nn.Linear(layer_sizes[-2], layer_sizes[-1]))

        # self.sequential = torch.nn.Sequential(*self.layers)

        self.init_weights()

    def init_weights(self):
        for layer in self.layers:
            if isinstance(layer, torch.nn.Linear):
                torch.nn.init.xavier_normal_(layer.weight)
                layer.bias.data.fill_(0.0)

    def forward(self, x, train=True):
        for layer in self.layers:
            x = layer(x)
            if train and hasattr(self, 'dropout_layer'):
                x = self.dropout_layer(x)

        return x
    
    def predict(self, x, train=False):
        self.eval()
        with torch.no_grad():
            return self.forward(x, train)

def train_neural_network(model, inputs, outputs, optimizer, epochs=1000, lr_scheduler=None):
    model.train()
    for epoch in range(epochs):
        optimizer.zero_grad()
        predictions = model(inputs)
        loss = torch.mean(torch.square(predictions - outputs))
        loss.backward()
        optimizer.step()

        if lr_scheduler:
            lr_scheduler.step()

        if epoch % 100 == 0:
            print(f'Epoch {epoch}, Loss: {loss.item()}, Learning Rate: {optimizer.param_groups[0]["lr"]}')


def load_model(model_path):
    checkpoint = torch.load(model_path, map_location=DEVICE)
    model_config = checkpoint['model_config']
    model = NeuralNetwork(model_config['layer_sizes'], dropout_rate=model_config['dropout_rate'])
    model.load_state_dict(checkpoint['model_state_dict'])
    print(f"Model loaded from {model_path}")

    model.to(DEVICE)
    model.eval()
    return model

def inverse_design(ply_number, y_target, n_restarts=20, epochs=1000, use_lbfgs=False):
    model = load_model('./model_checkpoint.pth')

    data = Dataset()
    y_target_norm = data.normalize_output(y_target) # (A1, B1, C1, Stress)
    y_target_tensor = torch.tensor(y_target, dtype=torch.float32)
    input_mean = torch.tensor(data.input_mean)
    input_std = torch.tensor(data.input_std)
    output_mean = torch.tensor(data.output_mean)
    output_std = torch.tensor(data.output_std)
    
    weights = torch.tensor([1.0, 1.0, 1.0, 0.0], dtype=torch.float32)
    bounds = torch.tensor([[1., 1000.], [1., 1000.], [1., 1000.]], dtype=torch.float32) # Initial_Temp, Punch_Velocity, Cooling_Time
    best = {"loss": float('inf'), "input": None, "output": None}

    for restart in range(n_restarts):
        z = torch.randn(3, requires_grad=True)

        if use_lbfgs:
            optimizer = torch.optim.LBFGS([z], lr=0.1, max_iter=epochs, line_search_fn="strong_wolfe")
            steps = 1
        else:
            optimizer = torch.optim.Adam([z], lr=0.001)
            steps = epochs
        
        for step in range(steps):
            def closure():
                var =  bounds[:, 0] + (bounds[:, 1] - bounds[:, 0]) * torch.sigmoid(z)
                optimizer.zero_grad()
                input_raw = torch.cat([torch.tensor([ply_number]), var]).unsqueeze(0)
                input_norm = (input_raw - input_mean) / input_std
                output_pred = model(input_norm)
                output_pred = (output_pred * output_std) + output_mean
                loss = torch.sum(weights * (output_pred - y_target_tensor) ** 2)
                loss.backward()
                return loss

            if use_lbfgs:
                loss = optimizer.step(closure)
            else:
                loss = closure()
                optimizer.step()

            if (step + 1) % 200 == 0:
                print(f'Restart {restart + 1}, Step {step + 1}, Loss: {loss.item():.6f}, grad: {z.grad.norm().item():.6f}')
            
        with torch.no_grad():
            var =  bounds[:, 0] + (bounds[:, 1] - bounds[:, 0]) * torch.sigmoid(z)
            input_raw = torch.cat([torch.tensor([ply_number]), var])
            input_norm = (input_raw - input_mean) / input_std
            output_pred = model(input_norm)
            output_pred = data.denormalize_output(output_pred.numpy())
            final_loss = np.sum(weights.numpy() * (output_pred - y_target) ** 2).item()
            if final_loss < best["loss"]:
                best["loss"] = final_loss
                best["input"] = var.detach().cpu().numpy()
                best["output"] = output_pred

    return best


def inverse_model():
    set_seed(5324)
    dataset = Dataset(inverse=True)
    inputs, outputs = dataset.get_input(normalize=True), dataset.get_output(normalize=True)

    idx_train = np.random.choice(len(inputs), size=int(0.85 * len(inputs)), replace=False)
    idx_test = np.setdiff1d(np.arange(len(inputs)), idx_train)
    # idx_test = np.array([1, 14+1, 18+1, 20+1, 23+1])
    # idx_train = np.setdiff1d(np.arange(len(inputs)), idx_test)

    inputs_train = torch.tensor(inputs[idx_train], dtype=torch.float32).to(DEVICE)
    outputs_train = torch.tensor(outputs[idx_train], dtype=torch.float32).to(DEVICE)

    inputs_test = torch.tensor(inputs[idx_test], dtype=torch.float32).to(DEVICE)
    outputs_test = torch.tensor(outputs[idx_test], dtype=torch.float32).to(DEVICE)
    
    layer_sizes = [inputs.shape[1]] + [64] * 3 + [outputs.shape[1]]
    model = NeuralNetwork(layer_sizes, dropout_rate=0.05, activation=torch.nn.ReLU).to(DEVICE)
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
    lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=5000, gamma=0.9)

    # Create a proper dataset that keeps input-output pairs together
    train_dataset = torch.utils.data.TensorDataset(inputs_train, outputs_train)
    train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=16, shuffle=True)

    # Train the model
    epochs = 20000
    for epoch in range(epochs):
        model.train()
        for inputs_batch, outputs_batch in train_loader:
            inputs_batch = inputs_batch.to(DEVICE)
            outputs_batch = outputs_batch.to(DEVICE)
            optimizer.zero_grad()
            predictions = model(inputs_batch)
            loss = torch.mean(torch.square(predictions - outputs_batch))
            loss.backward()
            optimizer.step()

        if lr_scheduler:
            lr_scheduler.step()

        if epoch % 500 == 0:
            train_pred = model(inputs_train)
            train_loss = torch.mean(torch.square(train_pred - outputs_train))
            test_pred = model(inputs_test)
            test_loss = torch.mean(torch.square(test_pred - outputs_test))
            print(f'Epoch {epoch}, Train Loss: {train_loss.item():.6f}, Test Loss: {test_loss.item():.6f}')
            # print(f'Learning Rate: {optimizer.param_groups[0]["lr"]}')


    predictions = model.predict(inputs_test)
    test_loss = torch.mean(torch.square(predictions - outputs_test))
    print(f'Test Loss: {test_loss.item()}. Samples: {idx_test}')

    x = np.arange(0, len(idx_test))

    outputs_test = dataset.denormalize_output(outputs_test.cpu().numpy())
    predictions = dataset.denormalize_output(predictions.cpu().numpy())
    # for sample in outputs_test:
    #     print(f'Test samples: {sample}')
    plt.figure(figsize=(10, 6))
    plt.plot(x, outputs_test[:, 0], color='b', linestyle='--', label='True Initial Temp')
    plt.plot(x, predictions[:, 0], color='b', linestyle='-', label='Predicted Initial Temp')
    plt.plot(x, outputs_test[:, 1], color='r', linestyle='--', label='True Punch Velocity')
    plt.plot(x, predictions[:, 1], color='r', linestyle='-', label='Predicted Punch Velocity')
    plt.plot(x, outputs_test[:, 2], color='g', linestyle='--', label='True Cooling Time')
    plt.plot(x, predictions[:, 2], color='g', linestyle='-', label='Predicted Cooling Time')
    plt.gca().xaxis.set_major_locator(plt.MaxNLocator(integer=True))    
    plt.xlabel('Sample Index')
    plt.xticks(ticks=range(len(idx_test)),labels=idx_test + 1)
    plt.ylabel('Processing Parameters')
    plt.legend(loc='upper right')
    plt.savefig('inverse_design.png')

    # MSE
    mse = np.mean((predictions - outputs_test) ** 2, axis=0)
    print(f'Mean Squared Error for Initial Temp: {mse[0]:.6f}, Punch Velocity: {mse[1]:.6f}, Cooling Time: {mse[2]:.6f}')

    # R 2 score
    ss_ress = np.sum((outputs_test - predictions) ** 2, axis=0)
    ss_tots = np.sum((outputs_test - np.mean(outputs_test, axis=0)) ** 2, axis=0)
    r2_scores = 1 - ss_ress / ss_tots
    print(f'R² Score for Initial Temp: {r2_scores[0]:.6f}, Punch Velocity: {r2_scores[1]:.6f}, Cooling Time: {r2_scores[2]:.6f}')

    # Error

    # Save the model
    model_save_path = './model_inverse_ckpt.pth'
    model_config = {'layer_sizes': layer_sizes,
                    'dropout_rate': 0.05
                    }
    checkpoint = {
        'model_state_dict': model.state_dict(),
        'model_config': model_config
    }
    torch.save(checkpoint, model_save_path)


if __name__ == "__main__":
    # train the inverse model over springback data
    # inverse_model()

    # perform inverse design
    import time
    start_time = time.time()
    best = inverse_design(ply_number=2, y_target=np.array([0.89, 0.83, 0.12, 180.2]), n_restarts=5, epochs=100, use_lbfgs=True)
    end_time = time.time()
    time_elapsed = (end_time - start_time)  # in milliseconds
    print(f"Inverse design completed in {time_elapsed:.2f} seconds.")
    print("Best Input (Initial Temp, Punch Velocity, Cooling Time):", best["input"])
    print("Best Output (A1, B1, C1, Stress):", best["output"])