backend/ml_module/test_chart_blocks.py

"""
Test script to verify chart and visualization blocks work end-to-end
"""
import json
from ml_module.core.response_formatter import (
    FormattedResponse,
    visualization_block,
    simple_table_with_types,
    chart_block,
    metric_block,
    text_block,
    Severity
)


def test_visualization_block():
    """Test creating visualization blocks"""
    print("\n=== Testing Visualization Block ===")
    
    # Feature importance data
    data = [
        {"feature": "age", "importance": 0.35},
        {"feature": "income", "importance": 0.28},
        {"feature": "credit_score", "importance": 0.22},
        {"feature": "loan_amount", "importance": 0.15},
    ]
    
    viz = visualization_block(
        data,
        chart_type="bar",
        title="Feature Importance",
        subtitle="Top features affecting model predictions"
    )
    
    print("✓ Created visualization block")
    print(json.dumps(viz.model_dump(), indent=2))
    return viz


def test_table_with_types():
    """Test creating tables with type inference"""
    print("\n=== Testing Table with Type Inference ===")
    
    # Time series data
    data = [
        {"date": "2024-01-01", "price": 150.5, "volume": 1000000},
        {"date": "2024-01-02", "price": 152.3, "volume": 1200000},
        {"date": "2024-01-03", "price": 151.8, "volume": 950000},
    ]
    
    table = simple_table_with_types(
        data,
        caption="Stock Prices",
        block_id="stock_data"
    )
    
    print("✓ Created table with inferred types")
    print(json.dumps(table.model_dump(), indent=2))
    return table


def test_formatted_response_with_charts():
    """Test a complete formatted response with charts"""
    print("\n=== Testing Complete FormattedResponse with Charts ===")
    
    # Simulate evaluation results
    metrics_data = [
        {"metric": "Accuracy", "score": 0.89},
        {"metric": "Precision", "score": 0.87},
        {"metric": "Recall", "score": 0.91},
        {"metric": "F1 Score", "score": 0.89},
    ]
    
    feature_data = [
        {"feature": "transaction_amount", "importance": 0.42},
        {"feature": "account_age_days", "importance": 0.28},
        {"feature": "num_transactions", "importance": 0.18},
        {"feature": "avg_transaction_size", "importance": 0.12},
    ]
    
    response = FormattedResponse(
        blocks=[
            text_block(
                "Model evaluation completed successfully",
                severity=Severity.SUCCESS
            ),
            metric_block("Accuracy", 0.89),
            metric_block("Precision", 0.87),
            metric_block("Recall", 0.91),
            metric_block("F1 Score", 0.89),
            visualization_block(
                metrics_data,
                chart_type="bar",
                title="Model Performance Metrics",
                subtitle="Evaluation on test set",
                block_id="eval_metrics_chart"
            ),
            visualization_block(
                feature_data,
                chart_type="bar",
                title="Feature Importance",
                subtitle="Top features affecting predictions",
                block_id="feature_importance_chart"
            ),
            simple_table_with_types(
                feature_data,
                caption="Detailed feature importance scores",
                block_id="feature_table"
            ),
        ],
        summary="Evaluation complete with visualizations",
        correlation_id="eval_v1_20241031",
        done=True
    )
    
    print("✓ Created complete formatted response")
    output = response.model_dump(mode="json")
    print(json.dumps(output, indent=2))
    
    # Verify structure
    assert len(output["blocks"]) == 8, "Should have 8 blocks"
    assert output["blocks"][5]["type"] == "visualization", "Block 5 should be visualization"
    assert output["blocks"][6]["type"] == "visualization", "Block 6 should be visualization"
    assert output["blocks"][7]["type"] == "table", "Block 7 should be table"
    
    print("✓ All assertions passed")
    return response


def test_echarts_chart_block():
    """Test creating an ECharts specification block"""
    print("\n=== Testing ECharts ChartBlock ===")
    
    echarts_spec = {
        "title": {"text": "Sales Over Time"},
        "xAxis": {
            "type": "category",
            "data": ["Jan", "Feb", "Mar", "Apr", "May"]
        },
        "yAxis": {"type": "value"},
        "series": [{
            "data": [150, 230, 224, 218, 135],
            "type": "line",
            "smooth": True
        }]
    }
    
    chart = chart_block(
        chart_type="line",
        specification=echarts_spec,
        title="Monthly Sales",
        subtitle="2024 Q1-Q2"
    )
    
    print("✓ Created ECharts chart block")
    print(json.dumps(chart.model_dump(), indent=2))
    return chart


if __name__ == "__main__":
    print("=" * 60)
    print("Testing ML Module Chart and Visualization Blocks")
    print("=" * 60)
    
    try:
        test_visualization_block()
        test_table_with_types()
        test_formatted_response_with_charts()
        test_echarts_chart_block()
        
        print("\n" + "=" * 60)
        print("✅ ALL TESTS PASSED!")
        print("=" * 60)
        print("\nBackend is now ready to send chart data to the frontend!")
        print("The frontend will automatically detect and render:")
        print("  • VisualizationBlock → Interactive charts with toggle buttons")
        print("  • TableBlock (with dtype) → Auto-converted to charts")
        print("  • ChartBlock → Direct ECharts rendering")
        
    except Exception as e:
        print(f"\n❌ Test failed: {e}")
        import traceback
        traceback.print_exc()
        exit(1)