mcp_server.py

"""
MCP Server for Agricultural Data Analysis
Provides tools and resources for analyzing agricultural intervention data.
"""

import json
import logging
from typing import Any, Dict, List, Optional
from mcp.server import Server
from mcp.server.models import InitializationOptions
from mcp.server.stdio import stdio_server
from mcp.types import Resource, Tool, TextContent
import asyncio
import pandas as pd
from data_loader import AgriculturalDataLoader
from analysis_tools import AgriculturalAnalyzer
import plotly.io as pio


# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("agricultural-mcp-server")

# Initialize data components
data_loader = AgriculturalDataLoader()
analyzer = AgriculturalAnalyzer(data_loader)

# Create MCP server
server = Server("agricultural-analysis")


@server.list_resources()
async def list_resources() -> List[Resource]:
    """List available resources."""
    return [
        Resource(
            uri="agricultural://data/summary",
            name="Data Summary",
            mimeType="application/json",
            description="Summary of available agricultural intervention data"
        ),
        Resource(
            uri="agricultural://data/years",
            name="Available Years",
            mimeType="application/json", 
            description="List of years with available data"
        ),
        Resource(
            uri="agricultural://data/plots",
            name="Available Plots",
            mimeType="application/json",
            description="List of available plots/parcels"
        ),
        Resource(
            uri="agricultural://data/crops",
            name="Available Crops",
            mimeType="application/json",
            description="List of available crop types"
        ),
        Resource(
            uri="agricultural://analysis/weed-pressure",
            name="Weed Pressure Analysis",
            mimeType="application/json",
            description="Current weed pressure trends analysis"
        ),
        Resource(
            uri="agricultural://analysis/rotation-impact",
            name="Crop Rotation Impact",
            mimeType="application/json",
            description="Analysis of crop rotation impact on weed pressure"
        )
    ]


@server.read_resource()
async def read_resource(uri: str) -> str:
    """Read a specific resource."""
    try:
        if uri == "agricultural://data/summary":
            df = data_loader.load_all_files()
            summary = {
                "total_records": len(df),
                "date_range": {
                    "start": df['datedebut'].min().strftime('%Y-%m-%d') if df['datedebut'].min() else None,
                    "end": df['datedebut'].max().strftime('%Y-%m-%d') if df['datedebut'].max() else None
                },
                "unique_plots": df['plot_name'].nunique(),
                "unique_crops": df['crop_type'].nunique(),
                "herbicide_applications": len(df[df['is_herbicide'] == True]),
                "years_covered": sorted(df['year'].unique().tolist())
            }
            return json.dumps(summary, indent=2)
            
        elif uri == "agricultural://data/years":
            years = data_loader.get_years_available()
            return json.dumps({"available_years": years})
            
        elif uri == "agricultural://data/plots":
            plots = data_loader.get_plots_available()
            return json.dumps({"available_plots": plots})
            
        elif uri == "agricultural://data/crops":
            crops = data_loader.get_crops_available()
            return json.dumps({"available_crops": crops})
            
        elif uri == "agricultural://analysis/weed-pressure":
            trends = analyzer.analyze_weed_pressure_trends()
            # Convert DataFrames to dict for JSON serialization
            serializable_trends = {}
            for key, value in trends.items():
                if isinstance(value, pd.DataFrame):
                    serializable_trends[key] = value.to_dict('records')
                else:
                    serializable_trends[key] = value
            return json.dumps(serializable_trends, indent=2)
            
        elif uri == "agricultural://analysis/rotation-impact":
            rotation_impact = analyzer.analyze_crop_rotation_impact()
            return json.dumps(rotation_impact.to_dict('records'), indent=2)
            
        else:
            raise ValueError(f"Unknown resource: {uri}")
            
    except Exception as e:
        logger.error(f"Error reading resource {uri}: {e}")
        return json.dumps({"error": str(e)})


@server.list_tools()
async def list_tools() -> List[Tool]:
    """List available tools."""
    return [
        Tool(
            name="filter_data",
            description="Filter agricultural data by years, plots, crops, or intervention types",
            inputSchema={
                "type": "object",
                "properties": {
                    "years": {
                        "type": "array",
                        "items": {"type": "integer"},
                        "description": "List of years to filter (e.g., [2022, 2023, 2024])"
                    },
                    "plots": {
                        "type": "array", 
                        "items": {"type": "string"},
                        "description": "List of plot names to filter"
                    },
                    "crops": {
                        "type": "array",
                        "items": {"type": "string"},
                        "description": "List of crop types to filter"
                    },
                    "intervention_types": {
                        "type": "array",
                        "items": {"type": "string"},
                        "description": "List of intervention types to filter"
                    }
                }
            }
        ),
        Tool(
            name="analyze_weed_pressure",
            description="Analyze weed pressure trends based on herbicide usage (IFT)",
            inputSchema={
                "type": "object",
                "properties": {
                    "years": {
                        "type": "array",
                        "items": {"type": "integer"},
                        "description": "Years to analyze"
                    },
                    "plots": {
                        "type": "array",
                        "items": {"type": "string"},
                        "description": "Plots to analyze"
                    },
                    "include_visualization": {
                        "type": "boolean",
                        "description": "Whether to include visualization data",
                        "default": True
                    }
                }
            }
        ),
        Tool(
            name="predict_weed_pressure",
            description="Predict weed pressure for the next 3 years using machine learning",
            inputSchema={
                "type": "object",
                "properties": {
                    "target_years": {
                        "type": "array",
                        "items": {"type": "integer"},
                        "description": "Years to predict (default: [2025, 2026, 2027])",
                        "default": [2025, 2026, 2027]
                    },
                    "plots": {
                        "type": "array",
                        "items": {"type": "string"},
                        "description": "Specific plots to predict for (optional)"
                    }
                }
            }
        ),
        Tool(
            name="identify_suitable_plots",
            description="Identify plots suitable for sensitive crops (peas, beans) based on low weed pressure",
            inputSchema={
                "type": "object",
                "properties": {
                    "target_years": {
                        "type": "array",
                        "items": {"type": "integer"},
                        "description": "Years to evaluate (default: [2025, 2026, 2027])",
                        "default": [2025, 2026, 2027]
                    },
                    "max_ift_threshold": {
                        "type": "number",
                        "description": "Maximum IFT threshold for suitable plots (default: 1.0)",
                        "default": 1.0
                    }
                }
            }
        ),
        Tool(
            name="analyze_crop_rotation",
            description="Analyze the impact of crop rotation patterns on weed pressure",
            inputSchema={
                "type": "object",
                "properties": {}
            }
        ),
        Tool(
            name="analyze_herbicide_alternatives",
            description="Analyze herbicide usage patterns and identify most used products",
            inputSchema={
                "type": "object",
                "properties": {}
            }
        ),
        Tool(
            name="get_data_statistics",
            description="Get comprehensive statistics about the agricultural data",
            inputSchema={
                "type": "object",
                "properties": {
                    "years": {
                        "type": "array",
                        "items": {"type": "integer"},
                        "description": "Years to analyze (optional)"
                    },
                    "plots": {
                        "type": "array",
                        "items": {"type": "string"},
                        "description": "Plots to analyze (optional)"
                    }
                }
            }
        )
    ]


@server.call_tool()
async def call_tool(name: str, arguments: Dict[str, Any]) -> List[TextContent]:
    """Execute a tool call."""
    try:
        if name == "filter_data":
            df = data_loader.filter_data(
                years=arguments.get("years"),
                plots=arguments.get("plots"),
                crops=arguments.get("crops"),
                intervention_types=arguments.get("intervention_types")
            )
            
            result = {
                "filtered_records": len(df),
                "summary": {
                    "unique_plots": df['plot_name'].nunique(),
                    "unique_crops": df['crop_type'].nunique(),
                    "year_range": [int(df['year'].min()), int(df['year'].max())] if len(df) > 0 else [],
                    "herbicide_applications": len(df[df['is_herbicide'] == True])
                },
                "sample_data": df.head(10).to_dict('records') if len(df) > 0 else []
            }
            
            return [TextContent(
                type="text",
                text=json.dumps(result, indent=2, default=str)
            )]
            
        elif name == "analyze_weed_pressure":
            trends = analyzer.analyze_weed_pressure_trends(
                years=arguments.get("years"),
                plots=arguments.get("plots")
            )
            
            # Convert DataFrames to dict for JSON serialization
            serializable_trends = {}
            for key, value in trends.items():
                if isinstance(value, pd.DataFrame):
                    serializable_trends[key] = value.to_dict('records')
                else:
                    serializable_trends[key] = value
            
            # Include visualization if requested
            if arguments.get("include_visualization", True):
                try:
                    fig = analyzer.create_weed_pressure_visualization(
                        years=arguments.get("years"),
                        plots=arguments.get("plots")
                    )
                    # Convert plot to HTML
                    serializable_trends["visualization_html"] = pio.to_html(fig, include_plotlyjs=True)
                except Exception as e:
                    serializable_trends["visualization_error"] = str(e)
            
            return [TextContent(
                type="text",
                text=json.dumps(serializable_trends, indent=2, default=str)
            )]
            
        elif name == "predict_weed_pressure":
            predictions = analyzer.predict_weed_pressure(
                target_years=arguments.get("target_years", [2025, 2026, 2027]),
                plots=arguments.get("plots")
            )
            
            # Convert DataFrames to dict for JSON serialization
            serializable_predictions = {}
            for key, value in predictions.items():
                if key == "predictions":
                    serializable_predictions[key] = {}
                    for year, df in value.items():
                        serializable_predictions[key][year] = df.to_dict('records')
                elif isinstance(value, pd.DataFrame):
                    serializable_predictions[key] = value.to_dict('records')
                else:
                    serializable_predictions[key] = value
            
            return [TextContent(
                type="text",
                text=json.dumps(serializable_predictions, indent=2, default=str)
            )]
            
        elif name == "identify_suitable_plots":
            suitable_plots = analyzer.identify_suitable_plots_for_sensitive_crops(
                target_years=arguments.get("target_years", [2025, 2026, 2027]),
                max_ift_threshold=arguments.get("max_ift_threshold", 1.0)
            )
            
            return [TextContent(
                type="text",
                text=json.dumps(suitable_plots, indent=2)
            )]
            
        elif name == "analyze_crop_rotation":
            rotation_impact = analyzer.analyze_crop_rotation_impact()
            
            return [TextContent(
                type="text",
                text=json.dumps(rotation_impact.to_dict('records'), indent=2, default=str)
            )]
            
        elif name == "analyze_herbicide_alternatives":
            herbicide_analysis = analyzer.analyze_herbicide_alternatives()
            
            return [TextContent(
                type="text",
                text=json.dumps(herbicide_analysis.to_dict('records'), indent=2, default=str)
            )]
            
        elif name == "get_data_statistics":
            df = data_loader.filter_data(
                years=arguments.get("years"),
                plots=arguments.get("plots")
            )
            
            stats = {
                "general": {
                    "total_records": len(df),
                    "unique_plots": df['plot_name'].nunique(),
                    "unique_crops": df['crop_type'].nunique(),
                    "date_range": {
                        "start": df['datedebut'].min().strftime('%Y-%m-%d') if not df['datedebut'].isna().all() else None,
                        "end": df['datedebut'].max().strftime('%Y-%m-%d') if not df['datedebut'].isna().all() else None
                    }
                },
                "interventions": {
                    "total_herbicide": len(df[df['is_herbicide'] == True]),
                    "total_fungicide": len(df[df['is_fungicide'] == True]),
                    "total_insecticide": len(df[df['is_insecticide'] == True])
                },
                "top_crops": df['crop_type'].value_counts().head(10).to_dict(),
                "top_plots": df['plot_name'].value_counts().head(10).to_dict(),
                "yearly_distribution": df['year'].value_counts().sort_index().to_dict()
            }
            
            return [TextContent(
                type="text", 
                text=json.dumps(stats, indent=2, default=str)
            )]
            
        else:
            raise ValueError(f"Unknown tool: {name}")
            
    except Exception as e:
        logger.error(f"Error executing tool {name}: {e}")
        return [TextContent(
            type="text",
            text=json.dumps({"error": str(e)}, indent=2)
        )]


async def main():
    """Main function to run the MCP server."""
    logger.info("Starting Agricultural MCP Server...")
    
    # Initialize the server
    async with stdio_server() as (read_stream, write_stream):
        await server.run(
            read_stream,
            write_stream,
            InitializationOptions(
                server_name="agricultural-analysis",
                server_version="1.0.0",
                capabilities=server.get_capabilities()
            )
        )


if __name__ == "__main__":
    asyncio.run(main())