import pandas as pd
import matplotlib.pyplot as plt
import math
import json
from scipy.stats import norm

from .patient import Patient
from pathlib import Path


DATA_PATH = Path(__file__).parent / "data"


class Child(Patient):
    """
    Represents a child patient (2–20 years) and provides growth chart analysis
    and visualization based on CDC growth standards.

    This class extends the Patient class and adds functionality for:
    - Weight-for-age
    - Length-for-age
    - BMI-for-age

    Supports plotting percentile curves and calculating Z-scores and percentiles
    using the LMS method.
    """

    # Child 2 to 20 database
    df_wtage = pd.read_csv(DATA_PATH /'wtage.csv', dtype={'Sex': str})
    df_lenage = pd.read_csv(DATA_PATH /'lenage.csv', dtype={'Sex': str})
    df_bmiage = pd.read_csv(DATA_PATH /'bmiage.csv', dtype={'Sex': str})

    def __init__(self, sex: str, age: int, wt: float = None, length: float = None):
        """
        Initialize a Child object.

        At least one anthropometric measurement must be provided.

        Args:
            sex (str): Child sex.
            age (int): Age in months.
            wt (float, optional): Weight in kilograms.
            length (float, optional): Length/height in centimeters.

        Raises:
            ValueError: If both wt and length are None.
        """
        if wt is None and length is None:
            raise ValueError("At least wt or length must be provided")

        super().__init__(sex, age, wt, length)

    def plot_wtage(self, title=None, figsize=(12,7), show=True, ax=None):
        """
    Plot weight-for-age growth chart for the child.

    This method draws CDC percentile curves and overlays
    the patient's weight measurement.

    The returned figure can be customized, embedded in apps,
    or saved as an image.

    Args:
        title (str, optional):
            Custom title for the chart.

        figsize (tuple, optional):
            Figure size in inches (width, height).
            Default is (12, 7).

        show (bool, optional):
            Whether to display the plot using matplotlib.
            If False, the plot will not be shown.
            Default is True.

        ax (matplotlib.axes.Axes, optional):
            Existing matplotlib Axes object for embedding
            in custom layouts or applications.

    Returns:
        tuple:
            (fig, ax)

            fig : matplotlib.figure.Figure
                The created matplotlib figure.

            ax : matplotlib.axes.Axes
                The axes containing the plot.
    """

        df = self.df_wtage
        sex = self._normalize_sex()

        data = df[df['Sex'] == sex]

        if ax is None:
            fig, ax = plt.subplots(figsize=figsize)
        else:
            fig = ax.figure

        for i, col in enumerate(self.agecolumns[1:]):

            if i in self.highlight_indices:
                ax.plot(
                    data['Agemos'], data[col],
                    color='gray', linewidth=2,
                    label=col
                )
            else:
                ax.plot(
                    data['Agemos'], data[col],
                    color='lightgray', linewidth=1
                )

        if self.wt is not None:
            ax.scatter(
                self.age, self.wt,
                color='red', s=100,
                zorder=5, label="Patient"
            )

        ax.set_title(title or "Weight for Age")
        ax.set_xlabel("Age (months)")
        ax.set_ylabel("Weight (Kg)")
        ax.grid(True)
        ax.legend()

        if show:
            plt.show()

        return fig, ax


    def plot_lenage(self, title=None, figsize=(12,7), show=True, ax=None):
        """
    Plot length/height-for-age growth chart for the child.

    This method draws CDC percentile curves and overlays
    the patient's length or height measurement.

    The output figure can be customized, embedded in
    applications, or exported as an image file.

    Args:
        title (str, optional):
            Custom title for the chart.

        figsize (tuple, optional):
            Figure size in inches (width, height).
            Default is (12, 7).

        show (bool, optional):
            Whether to display the plot using matplotlib.
            If False, the plot will not be shown.
            Default is True.

        ax (matplotlib.axes.Axes, optional):
            Existing matplotlib Axes object for embedding
            in custom layouts or applications.

    Returns:
        tuple:
            (fig, ax)

            fig : matplotlib.figure.Figure
                The created matplotlib figure.

            ax : matplotlib.axes.Axes
                The axes containing the plot.
    """

        df = self.df_lenage
        sex = self._normalize_sex()

        data = df[df['Sex'] == sex]

        if ax is None:
            fig, ax = plt.subplots(figsize=figsize)
        else:
            fig = ax.figure

        for i, col in enumerate(self.agecolumns[1:]):

            if i in self.highlight_indices:
                ax.plot(
                    data['Agemos'], data[col],
                    color='gray', linewidth=2,
                    label=col
                )
            else:
                ax.plot(
                    data['Agemos'], data[col],
                    color='lightgray', linewidth=1
                )

        if self.length is not None:
            ax.scatter(
                self.age, self.length,
                color='red', s=100,
                zorder=5, label="Patient"
            )

        ax.set_title(title or "Length for Age")
        ax.set_xlabel("Age (months)")
        ax.set_ylabel("Length (cm)")
        ax.grid(True)
        ax.legend()

        if show:
            plt.show()

        return fig, ax


    def plot_bmiage(self, title=None, figsize=(12,7), show=True, ax=None):
        """
    Plot BMI-for-age growth chart for the child.

    This method draws CDC BMI percentile curves and overlays
    the patient's calculated BMI value.

    The resulting figure can be modified, embedded in apps,
    or saved to disk.

    Args:
        title (str, optional):
            Custom title for the chart.

        figsize (tuple, optional):
            Figure size in inches (width, height).
            Default is (12, 7).

        show (bool, optional):
            Whether to display the plot using matplotlib.
            If False, the plot will not be shown.
            Default is True.

        ax (matplotlib.axes.Axes, optional):
            Existing matplotlib Axes object for embedding
            in custom layouts or applications.

    Returns:
        tuple:
            (fig, ax)

            fig : matplotlib.figure.Figure
                The created matplotlib figure.

            ax : matplotlib.axes.Axes
                The axes containing the plot.
    """

        df = self.df_bmiage
        sex = self._normalize_sex()

        data = df[df['Sex'] == sex]

        if ax is None:
            fig, ax = plt.subplots(figsize=figsize)
        else:
            fig = ax.figure

        for i, col in enumerate(self.agecolumns[1:]):

            if i in self.highlight_indices:
                ax.plot(
                    data['Agemos'], data[col],
                    color='gray', linewidth=2,
                    label=col
                )
            else:
                ax.plot(
                    data['Agemos'], data[col],
                    color='lightgray', linewidth=1
                )

        if self.wt and self.length and self.length > 0:
            ax.scatter(
                self.age, self.bmi,
                color='red', s=100,
                zorder=5, label="Patient"
            )

        ax.set_title(title or "BMI for Age")
        ax.set_xlabel("Age (months)")
        ax.set_ylabel("BMI")
        ax.grid(True)
        ax.legend()

        if show:
            plt.show()

        return fig, ax


def analyze_wtage(self, return_json=False):
    """
    Analyze weight-for-age using LMS method.

    Args:
        return_json (bool): If True, returns a JSON string. Default False.

    Returns:
        dict or JSON: Z-score, percentile, and interpretation.
    """
    df = self.df_wtage
    sex = self._normalize_sex()
    df2 = df[df['Sex'] == sex]

    row = df2.iloc[(df2['Agemos'] - self.age).abs().argsort()[:1]]
    L, M, S = float(row['L'].values[0]), float(row['M'].values[0]), float(row['S'].values[0])

    if L != 0:
        z = ((self.wt / M) ** L - 1) / (L * S)
    else:
        z = math.log(self.wt / M) / S

    percentile = norm.cdf(z) * 100

    result = {
        "indicator": "Weight-for-Age",
        "value": round(self.wt, 2),
        "unit": None,
        "Age_Months": round(self.age, 2),
        "sex": "Male" if sex == "1" else "Female",

        "z_score": round(z, 2),
        "percentile": round(percentile, 1),

        "z_interpretation": self._interpret_z(z),
        "percentile_interpretation": self._interpret_percentile(percentile)
    }

    if return_json:
        import json
        return json.dumps(result)

    return result


def analyze_lenage(self, return_json=False):
    """
    Analyze length-for-age using LMS method.

    Args:
        return_json (bool): If True, returns a JSON string. Default False.

    Returns:
        dict or JSON: Z-score, percentile, and interpretation.
    """
    df = self.df_lenage
    sex = self._normalize_sex()
    df2 = df[df['Sex'] == sex]

    row = df2.iloc[(df2['Agemos'] - self.age).abs().argsort()[:1]]
    L, M, S = float(row['L'].values[0]), float(row['M'].values[0]), float(row['S'].values[0])

    if L != 0:
        z = ((self.length / M) ** L - 1) / (L * S)
    else:
        z = math.log(self.length / M) / S

    percentile = norm.cdf(z) * 100

    result = {
        "indicator": "Length-for-Age",
        "value": round(self.length, 2),
        "unit": None,
        "Age_Months": round(self.age, 2),
        "sex": "Male" if sex == "1" else "Female",

        "z_score": round(z, 2),
        "percentile": round(percentile, 1),

        "z_interpretation": self._interpret_z(z),
        "percentile_interpretation": self._interpret_percentile(percentile)
    }

    if return_json:
        import json
        return json.dumps(result)

    return result


def analyze_bmiage(self, return_json=False):
    """
    Analyze BMI-for-age using LMS method.

    Args:
        return_json (bool): If True, returns a JSON string. Default False.

    Returns:
        dict or JSON: Z-score, percentile, and interpretation.
    """
    df = self.df_bmiage
    sex = self._normalize_sex()
    df2 = df[df['Sex'] == sex]

    row = df2.iloc[(df2['Agemos'] - self.age).abs().argsort()[:1]]
    L, M, S = float(row['L'].values[0]), float(row['M'].values[0]), float(row['S'].values[0])

    if L != 0:
        z = ((self.bmi / M) ** L - 1) / (L * S)
    else:
        z = math.log(self.bmi / M) / S

    percentile = norm.cdf(z) * 100

    result = {
        "indicator": "BMI-for-Age",
        "value": round(self.bmi, 2),
        "unit": None,
        "Age_Months": round(self.age, 2),
        "sex": "Male" if sex == "1" else "Female",

        "z_score": round(z, 2),
        "percentile": round(percentile, 1),

        "z_interpretation": self._interpret_z(z),
        "percentile_interpretation": self._interpret_percentile(percentile)
    }

    if return_json:
        import json
        return json.dumps(result)

    return result

