"""
Model Persistence and Serialization

Author: RSK World
Website: https://rskworld.in
Email: help@rskworld.in
Phone: +91 93305 39277
"""

import pickle
import json
import numpy as np
import pandas as pd
from datetime import datetime
import os
import warnings
warnings.filterwarnings('ignore')


class ModelPersistence:
    """
    Model Persistence and Serialization Tools
    
    Author: RSK World
    Website: https://rskworld.in
    Email: help@rskworld.in
    Phone: +91 93305 39277
    """
    
    def __init__(self, model_dir='saved_models'):
        """
        Initialize model persistence
        
        Parameters:
        -----------
        model_dir : str
            Directory to save models
        """
        self.model_dir = model_dir
        if not os.path.exists(model_dir):
            os.makedirs(model_dir)
    
    def save_model(self, model, model_name, metadata=None):
        """
        Save model to disk
        
        Parameters:
        -----------
        model : object
            Model object to save
        model_name : str
            Name for the model
        metadata : dict
            Additional metadata to save
        """
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        filename = f"{model_name}_{timestamp}.pkl"
        filepath = os.path.join(self.model_dir, filename)
        
        # Prepare save data
        save_data = {
            'model': model,
            'model_name': model_name,
            'timestamp': timestamp,
            'metadata': metadata or {}
        }
        
        # Save model
        with open(filepath, 'wb') as f:
            pickle.dump(save_data, f)
        
        # Save metadata as JSON
        json_filename = filename.replace('.pkl', '_metadata.json')
        json_filepath = os.path.join(self.model_dir, json_filename)
        
        json_metadata = {
            'model_name': model_name,
            'timestamp': timestamp,
            'metadata': metadata or {}
        }
        
        with open(json_filepath, 'w') as f:
            json.dump(json_metadata, f, indent=2, default=str)
        
        print(f"Model saved to: {filepath}")
        print(f"Metadata saved to: {json_filepath}")
        
        return filepath
    
    def load_model(self, filepath):
        """
        Load model from disk
        
        Parameters:
        -----------
        filepath : str
            Path to saved model file
        """
        with open(filepath, 'rb') as f:
            save_data = pickle.load(f)
        
        print(f"Model loaded: {save_data['model_name']}")
        print(f"Saved on: {save_data['timestamp']}")
        if save_data['metadata']:
            print(f"Metadata: {save_data['metadata']}")
        
        return save_data['model'], save_data['metadata']
    
    def list_saved_models(self):
        """List all saved models"""
        models = []
        for filename in os.listdir(self.model_dir):
            if filename.endswith('.pkl'):
                filepath = os.path.join(self.model_dir, filename)
                file_size = os.path.getsize(filepath)
                mod_time = datetime.fromtimestamp(os.path.getmtime(filepath))
                
                models.append({
                    'filename': filename,
                    'filepath': filepath,
                    'size': file_size,
                    'modified': mod_time
                })
        
        if models:
            df = pd.DataFrame(models)
            print("Saved Models:")
            print("=" * 70)
            print(df.to_string(index=False))
        else:
            print("No saved models found.")
        
        return models
    
    def delete_model(self, filepath):
        """Delete saved model"""
        if os.path.exists(filepath):
            os.remove(filepath)
            # Also remove metadata JSON if exists
            json_filepath = filepath.replace('.pkl', '_metadata.json')
            if os.path.exists(json_filepath):
                os.remove(json_filepath)
            print(f"Model deleted: {filepath}")
        else:
            print(f"File not found: {filepath}")
    
    def export_model_summary(self, model, filepath):
        """
        Export model summary to text file
        
        Parameters:
        -----------
        model : object
            Model with summary() method
        filepath : str
            Output file path
        """
        import sys
        from io import StringIO
        
        # Capture summary output
        old_stdout = sys.stdout
        sys.stdout = summary_capture = StringIO()
        
        try:
            model.summary()
            summary_text = summary_capture.getvalue()
        finally:
            sys.stdout = old_stdout
        
        # Write to file
        with open(filepath, 'w') as f:
            f.write(summary_text)
        
        print(f"Model summary exported to: {filepath}")
    
    def save_predictions(self, predictions, filepath, index=None):
        """
        Save predictions to CSV
        
        Parameters:
        -----------
        predictions : array-like
            Predictions to save
        filepath : str
            Output file path
        index : array-like
            Index for predictions
        """
        df = pd.DataFrame({
            'predictions': predictions
        }, index=index)
        
        df.to_csv(filepath)
        print(f"Predictions saved to: {filepath}")


if __name__ == "__main__":
    # Example usage
    print("Model Persistence Example")
    print("=" * 70)
    
    from regression_analysis import LinearRegressionModel
    import numpy as np
    
    # Create and fit a model
    np.random.seed(42)
    X = np.random.randn(100, 3)
    y = 2 + 1.5 * X[:, 0] + 0.8 * X[:, 1] - 0.5 * X[:, 2] + np.random.randn(100) * 0.5
    
    model = LinearRegressionModel()
    model.fit(X, y)
    
    # Save model
    persistence = ModelPersistence()
    metadata = {
        'n_samples': len(X),
        'n_features': X.shape[1],
        'r_squared': model.results.rsquared
    }
    
    filepath = persistence.save_model(model, 'linear_regression', metadata)
    
    # List saved models
    persistence.list_saved_models()
    
    # Load model
    loaded_model, loaded_metadata = persistence.load_model(filepath)
    
    # Export summary
    persistence.export_model_summary(loaded_model, 'model_summary.txt')