"""
Statistics Tools for TSA Agent

Tools for explaining statistical concepts and interpreting analysis results.
"""

from typing import Any
import math
from claude_agent_sdk import tool
from .api_client import format_tool_response


@tool(
    "explain_concept",
    "Provide detailed explanations of statistical concepts used in meta-analysis and TSA.",
    {"concept": str}
)
async def explain_concept(args: dict[str, Any]) -> dict[str, Any]:
    """
    Explain a statistical concept.

    Args:
        concept: The concept to explain - "heterogeneity", "i_squared", "odds_ratio",
                "relative_risk", "random_effects", "fixed_effect", "forest_plot",
                "funnel_plot", "meta_regression", "publication_bias", etc.
    """
    concept = args.get("concept", "").lower().replace(" ", "_").replace("-", "_")

    explanations = {
        "heterogeneity": """
## Heterogeneity in Meta-Analysis

**Heterogeneity** refers to variability in study results beyond what we'd expect from sampling error alone.

### Sources of Heterogeneity
1. **Clinical:** Different populations, interventions, or outcomes
2. **Methodological:** Different study designs, risk of bias
3. **Statistical:** Different effect measures or analysis methods

### Why It Matters
- High heterogeneity suggests studies may not be measuring the same thing
- A single pooled estimate may not be meaningful
- May need subgroup analysis or meta-regression

### Key Metrics
- **Q statistic:** Tests if variance exceeds expected (chi-squared test)
- **I²:** Percentage of variability due to heterogeneity (not chance)
- **τ² (tau-squared):** Estimated variance between studies
- **τ (tau):** Standard deviation between studies

### What to Do About It
1. Investigate sources (subgroup analysis)
2. Use random-effects model
3. Report with appropriate uncertainty
4. Consider if pooling is appropriate
""",

        "i_squared": """
## I-squared (I²) Statistic

I² describes the **percentage of variability in effect estimates that is due to heterogeneity** rather than sampling error.

### Formula
```
I² = max(0, (Q - df) / Q × 100%)
```
Where Q is the heterogeneity chi-squared and df = number of studies - 1.

### Interpretation Thresholds
| I² Value | Heterogeneity Level | Implication |
|----------|---------------------|-------------|
| 0-25% | Low | Studies fairly consistent |
| 25-50% | Low to Moderate | Some variation |
| 50-75% | Moderate to High | Substantial variation |
| >75% | High | Consider if pooling appropriate |

### Important Caveats
- I² is a **relative** measure, not absolute
- Can be high even with clinically unimportant variation
- Imprecise with few studies
- 0% doesn't guarantee homogeneity

### Clinical Judgment
Don't rely on I² alone – consider:
- Clinical similarity of studies
- Direction of effects (all favor same direction?)
- Prediction interval (range of true effects)
""",

        "odds_ratio": """
## Odds Ratio (OR)

The odds ratio compares the **odds** of an event between two groups.

### Formula
```
OR = (a/c) / (b/d) = (a × d) / (b × c)
```
Where in a 2×2 table:
- a = events in treatment, b = events in control
- c = non-events in treatment, d = non-events in control

### Interpretation
| OR Value | Meaning |
|----------|---------|
| OR = 1 | No difference between groups |
| OR < 1 | Event less likely in treatment group (protective) |
| OR > 1 | Event more likely in treatment group (harmful) |

### Example
OR = 0.5 means the **odds** of the event in treatment are half those in control.

### When to Use
✅ Case-control studies (RR not calculable)
✅ Rare events (approximates RR)
✅ Logistic regression outputs

### Limitations
⚠️ Not intuitive (odds ≠ probability)
⚠️ Overstates relative risk for common events
⚠️ Cannot directly calculate NNT
""",

        "relative_risk": """
## Relative Risk (RR) / Risk Ratio

Relative risk compares the **probability** (risk) of an event between groups.

### Formula
```
RR = Risk in treatment / Risk in control
   = (a / (a+c)) / (b / (b+d))
```

### Interpretation
| RR Value | Meaning |
|----------|---------|
| RR = 1 | No difference (null effect) |
| RR < 1 | Lower risk in treatment (beneficial if event is bad) |
| RR > 1 | Higher risk in treatment |

### Example
RR = 0.7 means **70% of the control group risk** in treatment group.
Or: **30% relative risk reduction** (RRR = 1 - RR = 0.30).

### Advantages Over OR
✅ More intuitive interpretation
✅ Directly measures risk
✅ Can calculate NNT: NNT = 1 / (Control Risk × (1 - RR))

### Limitations
⚠️ Cannot use in case-control studies
⚠️ Depends on baseline risk
⚠️ Relative measure – doesn't show absolute impact
""",

        "random_effects": """
## Random Effects Model

The random effects model assumes the **true effect varies across studies**.

### Concept
- Each study estimates a different true effect
- These true effects come from a distribution
- We estimate the mean of this distribution

### When to Use
✅ Studies differ in populations, interventions, settings
✅ Heterogeneity expected or observed (I² > 25%)
✅ You want results generalizable beyond included studies

### Common Methods
1. **DerSimonian-Laird (DL):** Most common, uses method of moments
2. **REML:** Restricted maximum likelihood, less biased
3. **Sidik-Jonkman:** Better with few studies
4. **Bayesian:** Incorporates prior information

### Effect on Results
- Wider confidence intervals than fixed effect
- Gives more weight to smaller studies
- Often more conservative (less likely to be "significant")

### Important Note
Random effects doesn't "fix" heterogeneity – it incorporates it into uncertainty.
""",

        "fixed_effect": """
## Fixed Effect Model

The fixed effect model assumes there is **one true effect** that all studies estimate.

### Concept
- All studies estimate the same underlying truth
- Differences between studies are only due to sampling error
- We estimate this single common effect

### When to Use
✅ Studies are essentially identical (same protocol)
✅ Heterogeneity is low (I² < 25%)
✅ You only want inference about included studies

### Weighting
Studies weighted by **inverse variance** only:
```
Weight = 1 / SE²
```
Larger, more precise studies get more weight.

### Limitations
⚠️ Assumes no heterogeneity (rarely true)
⚠️ Confidence intervals too narrow if heterogeneity exists
⚠️ Results only apply to these specific studies

### Comparison to Random Effects
| Aspect | Fixed | Random |
|--------|-------|--------|
| True effect | One | Distribution |
| CI width | Narrower | Wider |
| Small study weight | Less | More |
| Generalizability | Limited | Broader |
""",

        "forest_plot": """
## Forest Plot

A forest plot is the **standard visual summary** of a meta-analysis.

### Components
1. **Study labels:** Usually author and year (left side)
2. **Point estimates:** Squares showing each study's effect
3. **Confidence intervals:** Horizontal lines through squares
4. **Square size:** Proportional to study weight
5. **Diamond:** Pooled effect estimate (bottom)
6. **Vertical line:** Line of no effect (OR=1, RR=1, or MD=0)

### Reading the Plot
- **Square left of line:** Favors treatment
- **Square right of line:** Favors control (or harm)
- **CI crossing line:** Not statistically significant
- **Diamond crossing line:** Pooled effect not significant

### What to Look For
1. **Direction:** Do most studies favor same direction?
2. **Precision:** Are CIs narrow (precise) or wide (imprecise)?
3. **Consistency:** Do effects cluster or scatter?
4. **Outliers:** Any studies very different from others?
5. **Diamond vs squares:** Does pooling change the message?
""",

        "funnel_plot": """
## Funnel Plot

A funnel plot helps assess **publication bias** and small-study effects.

### Construction
- **X-axis:** Effect estimate (OR, RR, MD)
- **Y-axis:** Precision (often 1/SE or sample size)
- **Each point:** One study
- **Center line:** Pooled effect estimate

### Interpretation

**Symmetrical funnel = No evidence of publication bias**
- Small studies scatter evenly around pooled effect
- Larger studies cluster near the top, close to pooled effect

**Asymmetrical funnel = Potential bias**
- **Missing lower-left:** Unpublished negative small studies
- **Missing lower-right:** Unpublished positive small studies

### Caveats
⚠️ Need ~10 studies for reliable assessment
⚠️ Asymmetry can have other causes:
   - True heterogeneity
   - Different study populations
   - Methodological differences
   - Chance

### Statistical Tests
- **Egger's test:** Regression of effect on precision
- **Begg's test:** Rank correlation method
- **Trim and fill:** Imputes missing studies
""",

        "publication_bias": """
## Publication Bias

Publication bias occurs when **studies with certain results are more likely to be published**.

### The Problem
- Studies with "positive" (significant) results more likely published
- Studies with "negative" (null) results often unpublished
- Meta-analysis of published studies overestimates true effect

### Evidence of Publication Bias
1. **Funnel plot asymmetry:** Missing small negative studies
2. **Statistical tests:** Egger's, Begg's tests
3. **Excess of significant findings:** More p < 0.05 than expected
4. **Time-lag bias:** Positive results published faster

### Impact on TSA
- Inflates pooled effect estimate
- May lead to premature boundary crossing
- OIS calculation based on inflated effect

### Mitigation Strategies
1. **Search comprehensively:** Grey literature, trial registries
2. **Contact authors:** Request unpublished data
3. **Sensitivity analysis:** What if missing studies exist?
4. **Report transparently:** Acknowledge limitation
"""
    }

    if concept in explanations:
        return format_tool_response(explanations[concept])
    else:
        available = ", ".join(sorted(explanations.keys()))
        return format_tool_response(
            f"Concept '{concept}' not found.\n\n"
            f"**Available concepts:**\n{available}\n\n"
            f"Try one of these, or ask me to explain in your own words!"
        )


@tool(
    "interpret_heterogeneity",
    "Interpret heterogeneity statistics (I², Q, tau) from the current analysis.",
    {"i_squared": float, "q_statistic": float, "q_pvalue": float, "tau": float}
)
async def interpret_heterogeneity(args: dict[str, Any]) -> dict[str, Any]:
    """
    Interpret heterogeneity statistics.

    Args:
        i_squared: I² percentage (0-100)
        q_statistic: Cochran's Q value
        q_pvalue: P-value for Q test
        tau: Tau (between-study SD)
    """
    i2 = args.get("i_squared", 0)
    q = args.get("q_statistic", 0)
    q_p = args.get("q_pvalue", 1)
    tau = args.get("tau", 0)

    # Determine heterogeneity level
    if i2 < 25:
        level = "**Low**"
        level_desc = "Studies are fairly consistent. Fixed effect model may be appropriate."
        color = "🟢"
    elif i2 < 50:
        level = "**Low to Moderate**"
        level_desc = "Some variation exists. Random effects recommended."
        color = "🟡"
    elif i2 < 75:
        level = "**Moderate to High**"
        level_desc = "Substantial variation. Investigate sources with subgroup analysis."
        color = "🟠"
    else:
        level = "**High**"
        level_desc = "Consider if pooling is meaningful. Meta-regression may help."
        color = "🔴"

    # Q-test interpretation
    if q_p < 0.10:
        q_interp = "Q-test is **significant** (p < 0.10), confirming heterogeneity."
    else:
        q_interp = "Q-test is **not significant**, but has low power with few studies."

    summary = f"""
## Heterogeneity Interpretation

### Summary
{color} **Heterogeneity Level:** {level}

{level_desc}

### Statistics Breakdown

| Metric | Value | Interpretation |
|--------|-------|----------------|
| **I²** | {i2:.1f}% | {i2:.1f}% of variance due to heterogeneity |
| **Q** | {q:.2f} | Chi-squared test for heterogeneity |
| **Q p-value** | {q_p:.4f} | {q_interp} |
| **τ (tau)** | {tau:.4f} | Between-study standard deviation |

### Recommendations

1. **Model Choice:**
   {"Fixed effect may be considered" if i2 < 25 else "Use random effects model"}

2. **Next Steps:**
   {"Proceed with analysis" if i2 < 50 else "Investigate heterogeneity sources before interpreting pooled effect"}

3. **Reporting:**
   Always report I² and consider prediction interval for clinical interpretation.

### Prediction Interval
If I² > 0, the 95% prediction interval shows the range where 95% of true effects likely lie.
This is often more clinically meaningful than the confidence interval.
"""
    return format_tool_response(summary)


@tool(
    "interpret_effect",
    "Interpret a pooled effect estimate with its confidence interval.",
    {"effect": float, "ci_lower": float, "ci_upper": float, "measure": str, "z_score": float}
)
async def interpret_effect(args: dict[str, Any]) -> dict[str, Any]:
    """
    Interpret a pooled effect estimate.

    Args:
        effect: Point estimate (OR, RR, RD, or MD)
        ci_lower: Lower confidence interval bound
        ci_upper: Upper confidence interval bound
        measure: "OR" (Odds Ratio), "RR" (Relative Risk), "RD" (Risk Difference), "MD" (Mean Difference)
        z_score: Z-score for the effect
    """
    effect = args.get("effect", 1.0)
    ci_lower = args.get("ci_lower", 1.0)
    ci_upper = args.get("ci_upper", 1.0)
    measure = args.get("measure", "OR").upper()
    z = args.get("z_score", 0)

    # Determine null value and direction interpretation
    if measure in ("OR", "RR"):
        null = 1.0
        if effect < 1:
            direction = "favors treatment (reduces event rate)"
            magnitude = f"{(1 - effect) * 100:.1f}% relative reduction"
        elif effect > 1:
            direction = "favors control (increases event rate)"
            magnitude = f"{(effect - 1) * 100:.1f}% relative increase"
        else:
            direction = "no difference"
            magnitude = "0% change"
    else:  # RD or MD
        null = 0.0
        if effect < 0:
            direction = "favors treatment"
            magnitude = f"reduction of {abs(effect):.2f}"
        elif effect > 0:
            direction = "favors control"
            magnitude = f"increase of {effect:.2f}"
        else:
            direction = "no difference"
            magnitude = "no change"

    # Statistical significance
    ci_crosses_null = (ci_lower <= null <= ci_upper)
    if ci_crosses_null:
        sig_text = "**Not statistically significant** (95% CI crosses null)"
        sig_emoji = "⚠️"
    else:
        sig_text = "**Statistically significant** (95% CI excludes null)"
        sig_emoji = "✅" if ((measure in ("OR", "RR") and effect < 1) or (measure in ("RD", "MD") and effect < 0)) else "⚠️"

    # Calculate p-value approximation
    p_value = 2 * (1 - 0.5 * (1 + math.erf(abs(z) / math.sqrt(2)))) if z != 0 else 1.0

    summary = f"""
## Effect Estimate Interpretation

### Pooled Effect
| Metric | Value |
|--------|-------|
| **{measure}** | {effect:.3f} |
| **95% CI** | [{ci_lower:.3f}, {ci_upper:.3f}] |
| **Z-score** | {z:.3f} |
| **P-value** | {p_value:.4f} |

### Direction
The effect **{direction}**.
**Magnitude:** {magnitude}

### Statistical Significance
{sig_emoji} {sig_text}

### Clinical Interpretation Guidelines

{"**For Odds Ratio (OR):**" if measure == "OR" else ""}
{"- OR = " + f"{effect:.2f}" + " means odds in treatment are " + f"{effect:.0%}" + " of control odds" if measure == "OR" else ""}
{"- For rare events, OR ≈ RR" if measure == "OR" else ""}

{"**For Relative Risk (RR):**" if measure == "RR" else ""}
{"- RR = " + f"{effect:.2f}" + " means " + f"{effect:.0%}" + " of control group risk" if measure == "RR" else ""}
{"- Absolute risk reduction depends on baseline risk" if measure == "RR" else ""}

### Caveats
1. Statistical significance ≠ clinical importance
2. Consider effect size, not just p-value
3. Check heterogeneity before trusting pooled estimate
4. In TSA: check if boundaries crossed before concluding
"""
    return format_tool_response(summary)


@tool(
    "calculate_sample_size",
    "Calculate required sample size for detecting an effect in meta-analysis (OIS calculation).",
    {"p_ctrl": float, "p_int": float, "alpha": float, "power": float, "i_squared": float}
)
async def calculate_sample_size(args: dict[str, Any]) -> dict[str, Any]:
    """
    Calculate Optimal Information Size (required sample size).

    Args:
        p_ctrl: Expected control group event rate (0-1)
        p_int: Expected intervention group event rate (0-1)
        alpha: Type I error rate (typically 0.05)
        power: Desired power (typically 0.80)
        i_squared: Expected or observed I² (0-100) for heterogeneity adjustment
    """
    p_ctrl = args.get("p_ctrl", 0.15)
    p_int = args.get("p_int", 0.10)
    alpha = args.get("alpha", 0.05)
    power = args.get("power", 0.80)
    i2 = args.get("i_squared", 0)

    # Validate
    if not 0 < p_ctrl < 1 or not 0 < p_int < 1:
        return format_tool_response(
            "Error: Event rates must be between 0 and 1.",
            is_error=True
        )
    if p_ctrl == p_int:
        return format_tool_response(
            "Error: Control and intervention rates must differ.",
            is_error=True
        )

    # Calculate z-values
    z_alpha = abs(2.326 if alpha == 0.01 else 1.96 if alpha == 0.05 else 1.645)  # Two-sided
    z_beta = abs(0.842 if power == 0.80 else 1.282 if power == 0.90 else 1.645)

    # OIS formula for dichotomous outcomes
    p_star = (p_int + p_ctrl) / 2
    ois_unadjusted = 4 * (z_alpha + z_beta)**2 * (p_star * (1 - p_star)) / (p_ctrl - p_int)**2

    # Heterogeneity adjustment
    if i2 > 0 and i2 < 100:
        het_factor = 1 / (1 - i2/100)
        ois_adjusted = ois_unadjusted * het_factor
    else:
        het_factor = 1.0
        ois_adjusted = ois_unadjusted

    # Calculate relative risk reduction
    rrr = (p_ctrl - p_int) / p_ctrl * 100

    summary = f"""
## Sample Size Calculation (OIS)

### Input Parameters
| Parameter | Value |
|-----------|-------|
| Control event rate | {p_ctrl:.1%} |
| Intervention event rate | {p_int:.1%} |
| Relative Risk Reduction | {rrr:.1f}% |
| Alpha (Type I error) | {alpha} |
| Power (1 - Beta) | {power:.0%} |
| Heterogeneity (I²) | {i2:.0f}% |

### Results

**Unadjusted OIS:** {int(ois_unadjusted):,} patients
**Heterogeneity Factor:** {het_factor:.2f}
**Adjusted OIS:** {int(ois_adjusted):,} patients ⭐

### Interpretation

To reliably detect a {rrr:.0f}% relative risk reduction (from {p_ctrl:.1%} to {p_int:.1%}) with {power:.0%} power at α = {alpha}:

📊 You need approximately **{int(ois_adjusted):,} patients** in your meta-analysis.

### Notes
1. This is the **total** across treatment and control groups
2. OIS increases with higher heterogeneity (I²)
3. OIS increases when detecting smaller effects
4. This does NOT account for TSA repeated analyses – actual boundaries may require even more
"""
    return format_tool_response(summary)