import os, json, pickle, math, csv
import numpy as np
import matplotlib.pyplot as plt
from .mlp import MLP

class NeuralNetworkTrainer:
    
    def __init__(self, xs, ys, output_dir='output'):
        self.xs_original = xs
        self.ys_original = ys
        self.output_dir = output_dir
        
        # Create output directory
        os.makedirs(output_dir, exist_ok=True)
        
        # Validate data
        self._validate_and_prepare_data()
        
        # Default parameters (can be changed by user)
        self.params = {
            'architecture': [4, 4, 1],
            'normalization_method': 'max',
            'lr_finder_steps': 1000,
            'training_epochs': 3000,
            'model_filename': 'model.pkl',
            'norm_filename': 'norm_constants.json',
            'weights_filename': 'model_weights.json',
            'show_plots': True
        }
        
        # Create model
        self.model = MLP(self.input_dim, self.params['architecture'])
        
    def _validate_and_prepare_data(self):
        if len(self.xs_original) != len(self.ys_original):
            raise ValueError(f"Error: xs ({len(self.xs_original)}) and ys ({len(self.ys_original)}) must have same length")
        
        # Convert to numpy arrays
        xs_array = np.array(self.xs_original, dtype=float)
        ys_array = np.array(self.ys_original, dtype=float)
        
        # Get input dimension
        if xs_array.ndim == 1:
            self.input_dim = 1
            xs_array = xs_array.reshape(-1, 1)
        else:
            self.input_dim = xs_array.shape[1]
        
        # Flatten ys if needed
        if ys_array.ndim > 1:
            ys_array = ys_array[:, 0]
        
        self.xs_array = xs_array
        self.ys_array = ys_array
    
    def set_parameter(self, param_name, value):
        if param_name in self.params:
            self.params[param_name] = value
        else:
            raise ValueError(f"Unknown parameter: {param_name}")
    
    def get_parameter(self, param_name):
        return self.params.get(param_name)
    
    def normalize_data(self):
        method = self.params['normalization_method']
        
        if method == 'max':
            x_max = np.max(np.abs(self.xs_array))
            y_max = np.max(np.abs(self.ys_array))
            
            self.xs_norm = self.xs_array / x_max
            self.ys_norm = self.ys_array / y_max
            
            self.norm_constants = {
                'method': 'max',
                'x_max': float(x_max),
                'y_max': float(y_max)
            }
            
        elif method == 'minmax':
            x_min = np.min(self.xs_array, axis=0)
            x_max = np.max(self.xs_array, axis=0)
            y_min = np.min(self.ys_array)
            y_max = np.max(self.ys_array)
            
            self.xs_norm = (self.xs_array - x_min) / (x_max - x_min + 1e-8)
            self.ys_norm = (self.ys_array - y_min) / (y_max - y_min + 1e-8)
            
            self.norm_constants = {
                'method': 'minmax',
                'x_min': x_min.tolist(),
                'x_max': x_max.tolist(),
                'y_min': float(y_min),
                'y_max': float(y_max)
            }
        
        # Convert to list format for MLP
        self.xs = self.xs_norm.tolist()
        self.ys = self.ys_norm.tolist()
        
        return self.norm_constants
    
    def find_learning_rate(self, lr_start=1e-3, lr_end=1.0):
        self.model.reset_parameters()
        
        # Create learning rates on log scale
        steps = self.params['lr_finder_steps']
        lre = np.linspace(np.log10(lr_start), np.log10(lr_end), steps)
        lrs = 10 ** lre
        
        lr_list = []
        loss_list = []
        best_lr = lrs[0]
        best_loss = float('inf')
        
        for i in range(steps):
            # Forward pass
            ypred = [self.model(x) for x in self.xs]
            loss = sum((yout - ygt)**2 for ygt, yout in zip(self.ys, ypred))
            
            # Backward pass
            for p in self.model.parameters():
                p.grad = 0.0
            loss.backward()
            
            # Update with current learning rate
            lr = lrs[i]
            for p in self.model.parameters():
                p.data += -lr * p.grad
            
            # Store results
            lr_list.append(lr)
            loss_list.append(loss.data)
            
            # Update best LR
            if loss.data < best_loss:
                best_loss = loss.data
                best_lr = lr
        
        self.best_lr = best_lr
        self.lr_list = lr_list
        self.loss_list = loss_list
        
        return best_lr
    
    def train(self, learning_rate=None):
        if learning_rate is None:
            learning_rate = self.best_lr
        
        self.model.reset_parameters()
        
        epochs = self.params['training_epochs']
        losses, rmses, mapes, r2s = [], [], [], []
        actual_vs_predicted = []
        
        for epoch in range(epochs):
            # Forward pass
            ypred = [self.model(x) for x in self.xs]
            loss = sum((yout - ygt)**2 for ygt, yout in zip(self.ys, ypred))
            
            # Backward pass
            for p in self.model.parameters():
                p.grad = 0.0
            loss.backward()
            
            # Update weights
            for p in self.model.parameters():
                p.data += -learning_rate * p.grad
            
            # Calculate metrics
            y_true = [y for y in self.ys]
            y_pred = [y.data for y in ypred]
            
            mse = sum((yt - yp)**2 for yt, yp in zip(y_true, y_pred)) / len(self.ys)
            rmse = math.sqrt(mse)
            
            epsilon = 1e-8
            mape = sum(abs((yt - yp) / (yt + epsilon)) for yt, yp in zip(y_true, y_pred)) / len(self.ys)
            
            y_mean = sum(y_true) / len(y_true)
            ss_tot = sum((yt - y_mean)**2 for yt in y_true)
            ss_res = sum((yt - yp)**2 for yt, yp in zip(y_true, y_pred))
            r2 = 1 - (ss_res / ss_tot) if ss_tot != 0 else 0
            
            # Store metrics
            losses.append(loss.data)
            rmses.append(rmse)
            mapes.append(mape)
            r2s.append(r2)
            
            # Store actual vs predicted for last epoch
            if epoch == epochs - 1:
                for x_orig, y_orig, y_pred_norm in zip(self.xs_original, self.ys_original, ypred):
                    if self.norm_constants['method'] == 'max':
                        y_pred_denorm = y_pred_norm.data * self.norm_constants['y_max']
                    else:
                        y_pred_denorm = y_pred_norm.data * (self.norm_constants['y_max'] - self.norm_constants['y_min']) + self.norm_constants['y_min']
                    
                    actual_vs_predicted.append({
                        'input': x_orig,
                        'actual': y_orig,
                        'predicted': y_pred_denorm
                    })
        
        self.training_history = {
            'losses': losses,
            'rmses': rmses,
            'mapes': mapes,
            'r2s': r2s,
            'learning_rate': learning_rate,
            'final_loss': losses[-1],
            'final_r2': r2s[-1]
        }
        
        self.actual_vs_predicted = actual_vs_predicted
        
        return self.training_history
    
    def save_model(self):
        # Prepare model data
        params = [p.data for p in self.model.parameters()]
        
        model_data = {
            'model_params': params,
            'architecture': self.params['architecture'],
            'input_dim': self.input_dim,
            'training_history': self.training_history
        }
        
        model_path = os.path.join(self.output_dir, self.params['model_filename'])
        with open(model_path, 'wb') as f:
            pickle.dump(model_data, f)
        
        norm_path = os.path.join(self.output_dir, self.params['norm_filename'])
        with open(norm_path, 'w') as f:
            json.dump(self.norm_constants, f, indent=2)
        
        weights_path = os.path.join(self.output_dir, self.params['weights_filename'])
        self.save_weights_json(weights_path)
        
        csv_path = os.path.join(self.output_dir, 'actual_vs_predicted.csv')
        self.save_actual_vs_predicted_csv(csv_path)
        
        self._generate_plots()
        
        return model_path, norm_path, weights_path
    
    def save_weights_json(self, filepath):
        layers_data = self.model.get_all_weights_and_biases()
        
        json_data = {
            "model_info": {
                "input_dim": self.input_dim,
                "architecture": self.params['architecture'],
                "normalization": self.norm_constants,
                "training_info": {
                    "learning_rate": self.best_lr,
                    "final_loss": self.training_history['final_loss'],
                    "final_r2": self.training_history['final_r2']
                }
            },
            "layers": layers_data
        }
        
        with open(filepath, 'w') as f:
            json.dump(json_data, f, indent=2)
    
    def save_actual_vs_predicted_csv(self, filepath):
        with open(filepath, 'w', newline='') as f:
            writer = csv.writer(f)
            writer.writerow(['Input', 'Actual', 'Predicted', 'Error', 'Error_Percent'])
            
            for item in self.actual_vs_predicted:
                input_str = str(item['input'][0]) if len(item['input']) == 1 else str(item['input'])
                actual = item['actual']
                predicted = item['predicted']
                error = actual - predicted
                error_percent = (abs(error) / (abs(actual) + 1e-8)) * 100
                
                writer.writerow([input_str, actual, predicted, error, error_percent])
    
    def _generate_plots(self):
        if hasattr(self, 'lr_list') and hasattr(self, 'loss_list'):
            plt.figure(figsize=(10, 6))
            best_idx = np.argmin(self.loss_list)
            best_lr = self.lr_list[best_idx]
            best_loss = self.loss_list[best_idx]
            
            plt.plot(self.lr_list, self.loss_list, 'b-', linewidth=2)
            plt.scatter(best_lr, best_loss, color='red', s=100)
            plt.xscale('log')
            plt.xlabel('Learning Rate')
            plt.ylabel('Loss')
            plt.title('Learning Rate Finder')
            plt.grid(True, alpha=0.3)
            plt.tight_layout()
            plt.savefig(os.path.join(self.output_dir, 'learning_rate_finder.png'), dpi=150)
            plt.close()
        
        if hasattr(self, 'training_history'):
            epochs = list(range(self.params['training_epochs']))
            fig, axes = plt.subplots(2, 2, figsize=(14, 10))
            
            axes[0, 0].plot(epochs, self.training_history['losses'], color='blue', linewidth=2)
            axes[0, 0].set_xlabel('Epoch')
            axes[0, 0].set_ylabel('Loss')
            axes[0, 0].set_title('Training Loss')
            axes[0, 0].grid(True, alpha=0.3)
            
            axes[0, 1].plot(epochs, self.training_history['rmses'], color='green', linewidth=2)
            axes[0, 1].set_xlabel('Epoch')
            axes[0, 1].set_ylabel('RMSE')
            axes[0, 1].set_title('Root Mean Square Error')
            axes[0, 1].grid(True, alpha=0.3)
            
            axes[1, 0].plot(epochs, self.training_history['mapes'], color='orange', linewidth=2)
            axes[1, 0].set_xlabel('Epoch')
            axes[1, 0].set_ylabel('MAPE')
            axes[1, 0].set_title('Mean Absolute Percentage Error')
            axes[1, 0].grid(True, alpha=0.3)
            
            axes[1, 1].plot(epochs, self.training_history['r2s'], color='red', linewidth=2)
            axes[1, 1].set_xlabel('Epoch')
            axes[1, 1].set_ylabel('R2 Score')
            axes[1, 1].set_title('R2 Score')
            axes[1, 1].grid(True, alpha=0.3)
            
            plt.tight_layout()
            plt.savefig(os.path.join(self.output_dir, 'training_metrics.png'), dpi=150)
            plt.close()
        
        if hasattr(self, 'actual_vs_predicted'):
            plt.figure(figsize=(10, 6))
            actuals = [item['actual'] for item in self.actual_vs_predicted]
            predictions = [item['predicted'] for item in self.actual_vs_predicted]
            
            plt.scatter(actuals, predictions, color='blue', alpha=0.6)
            
            # Add perfect prediction line
            min_val = min(min(actuals), min(predictions))
            max_val = max(max(actuals), max(predictions))
            plt.plot([min_val, max_val], [min_val, max_val], 'r--', label='Perfect Prediction')
            
            plt.xlabel('Actual Values')
            plt.ylabel('Predicted Values')
            plt.title('Actual vs Predicted Values')
            plt.grid(True, alpha=0.3)
            plt.legend()
            plt.tight_layout()
            plt.savefig(os.path.join(self.output_dir, 'actual_vs_predicted.png'), dpi=150)
            plt.close()
    
    def run_training(self):
        print("Starting neural network training...")
        print(f"Data points: {len(self.xs_original)}")
        print(f"Architecture: MLP({self.input_dim}, {self.params['architecture']})")
        print(f"Output directory: {self.output_dir}")
        print("-" * 50)
        
        print("Step 1: Normalizing data...")
        self.normalize_data()
        
        print("Step 2: Finding optimal learning rate...")
        best_lr = self.find_learning_rate()
        print(f"  Best learning rate: {best_lr:.6f}")
        
        print(f"Step 3: Training model ({self.params['training_epochs']} epochs)...")
        history = self.train(learning_rate=best_lr)
        print(f"  Final loss: {history['final_loss']:.6f}")
        print(f"  Final R2: {history['final_r2']:.6f}")
        
        print("Step 4: Saving model and results...")
        model_path, norm_path, weights_path = self.save_model()
        
        print("-" * 50)
        print("Training completed successfully!")
        print(f"Model saved: {model_path}")
        print(f"Normalization constants: {norm_path}")
        print(f"Weights JSON (for C++): {weights_path}")
        print(f"Output files in: {self.output_dir}")
        
        return model_path, norm_path, weights_path

