import pandas as pd
import numpy as np
from momentfm import MOMENTPipeline
from io import StringIO
import gradio as gr

# Initialize model with proper configuration
model = MOMENTPipeline.from_pretrained(
    "AutonLab/MOMENT-1-large",
    model_kwargs={"task_name": "reconstruction"},
)
model.init()

def generate_analysis_report(data_input, sensitivity=3.0):
    """Generate comprehensive textual analysis report"""
    try:
        # Process and validate input data
        df = pd.read_csv(StringIO(data_input))
        
        if 'timestamp' not in df.columns or 'value' not in df.columns:
            return "Error: CSV must contain 'timestamp' and 'value' columns"
            
        df['timestamp'] = pd.to_datetime(df['timestamp'], errors='coerce')
        df['value'] = pd.to_numeric(df['value'], errors='coerce')
        
        if df.isnull().values.any():
            return "Error: Invalid data in timestamp or value columns"
            
        df = df.sort_values('timestamp').dropna()
        
        # Prepare data for model (3D array format)
        values = df['value'].values.astype(np.float32)
        values_3d = values.reshape(1, -1, 1)  # Reshape to [batch, sequence, features]
        
        # Correct reconstruction call with proper parameter
        reconstructed = model.reconstruct(X=values_3d)  # Using named parameter
        
        # Calculate errors and detect anomalies
        errors = np.abs(values - reconstructed[0,:,0])
        median = np.median(errors)
        mad = np.median(np.abs(errors - median))
        threshold = median + sensitivity * (1.4826 * mad)
        
        # Identify anomalies
        anomalies = df.copy()
        anomalies['anomaly_score'] = errors
        anomalies = anomalies[errors > threshold].sort_values('anomaly_score', ascending=False)
        
        # Generate report
        report = f"""
EQUIPMENT ANALYSIS REPORT
========================
Generated: {pd.Timestamp.now().strftime('%Y-%m-%d %H:%M:%S')}
Sensitivity: {sensitivity} (z-score)

DATA SUMMARY
------------
Time period: {df['timestamp'].min()} to {df['timestamp'].max()}
Data points: {len(df)}
Value range: {df['value'].min():.2f} to {df['value'].max():.2f}
Median value: {df['value'].median():.2f}

ANOMALY FINDINGS
----------------
Detection threshold: {threshold:.2f}
Anomalies found: {len(anomalies)} ({len(anomalies)/len(df):.1%})
Most severe: {errors.max():.2f} at {df.loc[errors.argmax(), 'timestamp']}

TOP ANOMALIES
-------------
{anomalies[['timestamp', 'value', 'anomaly_score']].head(10).to_string(index=False, float_format='%.2f')}

RECOMMENDATIONS
---------------
1. Investigate top 3 anomalies for potential equipment issues
2. Check maintenance records around {anomalies['timestamp'].iloc[0].strftime('%Y-%m-%d %H:%M')}
3. Consider recalibration if anomalies persist
4. Review sensor health if anomalies cluster in time
"""
        return report.strip()
    
    except Exception as e:
        return f"ANALYSIS FAILED: {str(e)}"

# Gradio Interface
with gr.Blocks(title="Equipment Analysis Reporter") as demo:
    gr.Markdown("## 🏭 Equipment Health Analysis Report")
    
    with gr.Row():
        with gr.Column():
            data_input = gr.Textbox(
                label="Paste CSV Data (timestamp,value)",
                value="""timestamp,value
2025-04-01 00:00:00,100
2025-04-01 01:00:00,102
2025-04-01 02:00:00,98
2025-04-01 03:00:00,105
2025-04-01 04:00:00,103
2025-04-01 05:00:00,107
2025-04-01 06:00:00,200
2025-04-01 07:00:00,108
2025-04-01 08:00:00,110
2025-04-01 09:00:00,98
2025-04-01 10:00:00,99
2025-04-01 11:00:00,102
2025-04-01 12:00:00,101""",
                lines=10
            )
            sensitivity = gr.Slider(1.0, 5.0, value=3.0, step=0.1, label="Detection Sensitivity")
            submit_btn = gr.Button("Generate Report", variant="primary")
            
        with gr.Column():
            report_output = gr.Textbox(
                label="Analysis Report", 
                lines=20,
                interactive=False
            )
    
    submit_btn.click(
        generate_analysis_report,
        inputs=[data_input, sensitivity],
        outputs=report_output
    )

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)