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

from .patient import Patient
from pathlib import Path


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


class Infant(Patient):
    """
    Represents an infant patient 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
    - Head circumference-for-age
    - Weight-for-length

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

    # infant database
    df_wtageinf = pd.read_csv(DATA_PATH / 'wtageinf.csv', dtype={'Sex': str})
    df_lenageinf = pd.read_csv(DATA_PATH / 'lenageinf.csv', dtype={'Sex': str})
    df_hcageinf = pd.read_csv(DATA_PATH / 'hcageinf.csv', dtype={'Sex': str})
    df_wtleninf = pd.read_csv(DATA_PATH / 'wtleninf.csv', dtype={'Sex': str})

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

        At least one anthropometric measurement must be provided.

        Args:
            sex (str):
                Infant sex.

            age (int):
                Age in months.

            wt (float, optional):
                Weight in kilograms.

            length (float, optional):
                Length in centimeters.

            hc (float, optional):
                Head circumference in centimeters.

        Raises:
            ValueError:
                If all wt, length, and hc are None.
        """

        if wt is None and length is None and hc is None:
            raise ValueError("At least wt or length must be provided")

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

        self.hc = hc

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

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

    The resulting 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_wtageinf
        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 "Infant 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_lenageinf(self, title=None, figsize=(12,7), show=True, ax=None):
        """
    Plot length-for-age growth chart for infants.

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

    The generated figure can be customized, embedded
    in applications, 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_lenageinf
        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 "Infant 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_hcageinf(self, title=None, figsize=(12,7), show=True, ax=None):
        """
    Plot head circumference-for-age growth chart for infants.

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

    The output figure can be customized, embedded in apps,
    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_hcageinf
        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.hc is not None:
            ax.scatter(
                self.age, self.hc,
                color='red', s=100,
                zorder=5, label="Patient"
            )

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

        if show:
            plt.show()

        return fig, ax
    
    def plot_wtleninf(self, title=None, figsize=(12,7), show=True, ax=None):
        """
    Plot weight-for-length growth chart for infants.

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

    It is mainly used for nutritional and growth status
    assessment in infants.

    The generated 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_wtleninf
        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.lencolumns[1:]):

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

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

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

        if show:
            plt.show()

        return fig, ax



    def analyze_wtageinf(self, as_json=False):
        """
        Analyze weight-for-age using LMS method (WHO/CDC).

        Returns:
            dict or json:
                Z-score, percentile, and interpretation.
        """

        df = self.df_wtageinf

        sex = self._normalize_sex()

        df2 = df[df['Sex'] == sex]

        # Closest age row
        row = df2.iloc[
            (df2['Agemos'] - self.age).abs().argsort()[:1]
        ]

        L = float(row['L'].values[0])
        M = float(row['M'].values[0])
        S = float(row['S'].values[0])

        # LMS Z-score
        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": "kg",
            "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 as_json:
            return json.dumps(result, ensure_ascii=False)

        return result


def analyze_lenageinf(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_lenageinf
    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": "cm",
    "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_hcageinf(self, return_json=False):
    """
    Analyze head circumference-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_hcageinf
    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.hc / M) ** L - 1) / (L * S)
    else:
        z = math.log(self.hc / M) / S

    percentile = norm.cdf(z) * 100

    result = {
        "indicator": "Head Circumference-for-Age",
        "value": round(self.hc, 2),
        "unit": "cm",
        "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_wtleninf(self, return_json=False):
    """
    Analyze weight-for-length 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_wtleninf
    sex = self._normalize_sex()
    df2 = df[df['Sex'] == sex]

    row = df2.iloc[(df2['Length'] - self.length).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-Length",
        "value": round(self.wt, 2),
        "unit": None,
        "length_cm": round(self.length, 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
