agent/src/tools/meta_analysis_tool.py

"""
Meta-Analysis Tools for TSA Agent

Tools for creating and configuring meta-analyses using the TSA Java engine.
"""

from typing import Any
from claude_agent_sdk import tool
from .api_client import api_call, format_tool_response


@tool(
    "create_meta_analysis",
    "Create a new meta-analysis with specified settings. Returns the created analysis configuration.",
    {
        "name": str,
        "group1": str,
        "group2": str,
        "trial_type": str,
        "effect_model": str,
        "effect_measure": str,
    }
)
async def create_meta_analysis(args: dict[str, Any]) -> dict[str, Any]:
    """
    Create a new meta-analysis.

    Args:
        name: Analysis name (e.g., "Mortality Meta-Analysis")
        group1: Intervention/treatment group name
        group2: Control/comparison group name
        trial_type: "Dichotomous" or "Continuous"
        effect_model: "Fixed", "RandomDL", "RandomSJ", "RandomBT", "HybridDL", "HybridBT"
        effect_measure: "OddsRatio", "RelativeRisk", "RiskDifference", "MeanDifference"
    """
    # Validate trial type
    if args.get("trial_type", "").lower() not in ("dichotomous", "continuous"):
        return format_tool_response(
            "Invalid trial_type. Use 'Dichotomous' or 'Continuous'.",
            is_error=True
        )

    # Validate effect model
    valid_models = ["Fixed", "RandomDL", "RandomSJ", "RandomBT", "HybridDL", "HybridBT"]
    if args.get("effect_model", "") not in valid_models:
        return format_tool_response(
            f"Invalid effect_model. Use one of: {', '.join(valid_models)}",
            is_error=True
        )

    # Call API
    response = await api_call(
        "POST",
        "/analysis/create",
        params={
            "name": args.get("name", "New Analysis"),
            "group1": args.get("group1", "Treatment"),
            "group2": args.get("group2", "Control"),
            "trial_type": args.get("trial_type", "Dichotomous"),
            "effect_model": args.get("effect_model", "RandomDL"),
            "effect_measure": args.get("effect_measure", "OddsRatio"),
        }
    )

    if response.get("success"):
        summary = f"""
## Meta-Analysis Created Successfully

**Name:** {args.get("name", "New Analysis")}

**Configuration:**
- Trial Type: {args.get("trial_type", "Dichotomous")}
- Effect Model: {args.get("effect_model", "RandomDL")}
- Effect Measure: {args.get("effect_measure", "OddsRatio")}
- Groups: {args.get("group1", "Treatment")} vs {args.get("group2", "Control")}

You can now add trials using `add_dichotomous_trial` or `add_continuous_trial`.
"""
        return format_tool_response(summary)
    else:
        return format_tool_response(response)


@tool(
    "get_analysis_state",
    "Get the current state of the meta-analysis including whether one is loaded and basic info.",
    {}
)
async def get_analysis_state(args: dict[str, Any]) -> dict[str, Any]:
    """Check if a meta-analysis is currently loaded and get its state."""
    response = await api_call("GET", "/analysis/state")

    if response.get("has_analysis"):
        results = response.get("results", {})
        summary = f"""
## Current Analysis State

**Analysis Loaded:** Yes
**Trials:** {results.get("num_trials", 0)}
**Total Patients:** {results.get("total_patients", 0):,}
**Pooled Effect:** {results.get("pooled_effect", "N/A")}
**Boundaries Configured:** {response.get("num_boundaries", 0)}
"""
        return format_tool_response(summary)
    else:
        return format_tool_response(
            "No meta-analysis is currently loaded. "
            "Create one using `create_meta_analysis` or load from a .TSA file."
        )


@tool(
    "get_analysis_results",
    "Get detailed results from the current meta-analysis including pooled effect, heterogeneity, and confidence intervals.",
    {}
)
async def get_analysis_results(args: dict[str, Any]) -> dict[str, Any]:
    """Get comprehensive results from the current analysis."""
    # Get main results
    results_response = await api_call("GET", "/results")

    if not results_response.get("success"):
        return format_tool_response(results_response)

    results = results_response.get("results", {})

    # Get confidence intervals
    ci_response = await api_call("GET", "/results/confidence-intervals")
    intervals = ci_response.get("intervals", {}) if ci_response.get("success") else {}

    # Format nice output
    summary = f"""
## Meta-Analysis Results

### Pooled Effect Estimate
- **Effect:** {results.get("pooled_effect", "N/A")}
- **Z-score:** {results.get("z_score", "N/A")}
- **P-value:** {results.get("p_value", "N/A")}

### Confidence Intervals
| Level | Lower | Upper |
|-------|-------|-------|
| 90% | {intervals.get("90%", {}).get("lower", "N/A")} | {intervals.get("90%", {}).get("upper", "N/A")} |
| 95% | {intervals.get("95%", {}).get("lower", "N/A")} | {intervals.get("95%", {}).get("upper", "N/A")} |
| 99% | {intervals.get("99%", {}).get("lower", "N/A")} | {intervals.get("99%", {}).get("upper", "N/A")} |

### Heterogeneity
- **Q Statistic:** {results.get("heterogeneity_q", "N/A")}
- **I²:** {results.get("i_squared", "N/A")}%
- **Tau²:** {results.get("tau_squared", "N/A")}
- **Tau:** {results.get("tau", "N/A")}

### Sample Information
- **Number of Trials:** {results.get("num_trials", 0)}
- **Total Patients:** {results.get("total_patients", 0):,}
- **Total Events:** {results.get("total_events", 0):,}
"""
    return format_tool_response(summary)


@tool(
    "update_analysis_settings",
    "Update the meta-analysis settings like effect model, effect measure, or zero-event handling.",
    {
        "effect_model": str,
        "effect_measure": str,
        "zero_handling": str,
        "zero_value": float,
    }
)
async def update_analysis_settings(args: dict[str, Any]) -> dict[str, Any]:
    """
    Update meta-analysis configuration.

    Args:
        effect_model: "Fixed", "RandomDL", "RandomSJ", etc.
        effect_measure: "OddsRatio", "RelativeRisk", "RiskDifference", "MeanDifference"
        zero_handling: "Ignore", "Constant", "Empirical" (for handling zero events)
        zero_value: Continuity correction value (default 0.5)
    """
    # This would call an update endpoint
    # For now, return guidance
    summary = f"""
## Settings Update

To update analysis settings:

1. **Effect Model**: Controls how studies are combined
   - Fixed: Assumes single true effect
   - RandomDL: DerSimonian-Laird random effects
   - RandomSJ: Sidik-Jonkman estimator
   - RandomBT: Bayesian tau estimation
   - HybridDL/HybridBT: Hybrid approaches

2. **Effect Measure**: Type of effect size
   - OddsRatio: Good for rare events
   - RelativeRisk: More intuitive interpretation
   - RiskDifference: Absolute effect size
   - MeanDifference: For continuous outcomes

3. **Zero Event Handling**: How to handle studies with no events
   - Ignore: Exclude zero-event studies
   - Constant: Add continuity correction (typically 0.5)
   - Empirical: Proportion-based correction

Settings requested:
- Effect Model: {args.get("effect_model", "Not specified")}
- Effect Measure: {args.get("effect_measure", "Not specified")}
- Zero Handling: {args.get("zero_handling", "Not specified")}
- Zero Value: {args.get("zero_value", "Not specified")}

*Note: Settings are applied when recalculating results.*
"""
    return format_tool_response(summary)