import pandas as pd
import matplotlib.pyplot as plt
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):
        """
        Plot weight-for-age growth chart.

        Displays percentile curves and overlays patient's weight.

        Args:
            title (str, optional):
                Title of the plot.
        """

        df = self.df_wtageinf

        self.sex = str(self.sex).lower()

        if self.sex in ['m', '1']:
            self.sex = '1'
        elif self.sex in ['f', '2']:
            self.sex = '2'
        else:
            raise TypeError('define sex as m or f')

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

        plt.figure(figsize=(12, 7))

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

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

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

        plt.title(title)
        plt.xlabel("age(month)")
        plt.ylabel("Weight(Kg)")
        plt.grid(True)
        plt.legend()
        plt.show()

    def plot_lenageinf(self, title=None):
        """
        Plot length-for-age growth chart.

        Displays percentile curves and overlays patient's length.

        Args:
            title (str, optional):
                Title of the plot.
        """

        df = self.df_lenageinf

        self.sex = str(self.sex).lower()

        if self.sex in ['m', '1']:
            self.sex = '1'
        elif self.sex in ['f', '2']:
            self.sex = '2'
        else:
            raise TypeError('define sex as m or f')

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

        plt.figure(figsize=(14, 10))

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

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

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

        plt.title(title)
        plt.xlabel("age(month)")
        plt.ylabel("Length(cm)")
        plt.grid(True)
        plt.legend()
        plt.show()

    def plot_hcageinf(self, title=None):
        """
        Plot head circumference-for-age growth chart.

        Displays percentile curves and overlays patient's head circumference.

        Args:
            title (str, optional):
                Title of the plot.
        """

        df = self.df_hcageinf

        self.sex = str(self.sex).lower()

        if self.sex in ['m', '1']:
            self.sex = '1'
        elif self.sex in ['f', '2']:
            self.sex = '2'
        else:
            raise TypeError('define sex as m or f')

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

        plt.figure(figsize=(12, 7))

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

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

        if self.hc is not None:
            plt.scatter(self.age, self.hc,
                        color='red', s=100, zorder=5)

        plt.title(title)
        plt.xlabel("age(month)")
        plt.ylabel("Head Circumference(cm)")
        plt.grid(True)
        plt.legend()
        plt.show()

    def analyze_wtageinf(self):
        """
        Analyze weight-for-age using LMS method.

        Returns:
            dict:
                Z-score and percentile.
        """

        df = self.df_wtageinf

        self.sex = str(self.sex).lower()

        if self.sex in ['m', '1']:
            self.sex = '1'
        elif self.sex in ['f', '2']:
            self.sex = '2'
        else:
            raise TypeError('define sex as m or f')

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

        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])

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

        percentile = norm.cdf(z) * 100

        return {
            "Z": round(z, 2),
            "Percentile": round(percentile, 2)
        }

    def analyze_lenageinf(self):
        """
        Analyze length-for-age using LMS method.

        Returns:
            dict:
                Z-score and percentile.
        """

        df = self.df_lenageinf

        self.sex = str(self.sex).lower()

        if self.sex in ['m', '1']:
            self.sex = '1'
        elif self.sex in ['f', '2']:
            self.sex = '2'
        else:
            raise TypeError('define sex as m or f')

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

        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])

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

        percentile = norm.cdf(z) * 100

        return {
            "Z": round(z, 2),
            "Percentile": round(percentile, 2)
        }

    def analyze_hcageinf(self):
        """
        Analyze head circumference-for-age using LMS method.

        Returns:
            dict:
                Z-score and percentile.
        """

        df = self.df_hcageinf

        self.sex = str(self.sex).lower()

        if self.sex in ['m', '1']:
            self.sex = '1'
        elif self.sex in ['f', '2']:
            self.sex = '2'
        else:
            raise TypeError('define sex as m or f')

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

        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])

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

        percentile = norm.cdf(z) * 100

        return {
            "Z": round(z, 2),
            "Percentile": round(percentile, 2)
        }

    def analyze_wtleninf(self):
        """
        Analyze weight-for-length using LMS method.

        Returns:
            dict:
                Z-score and percentile.
        """

        df = self.df_wtleninf

        self.sex = str(self.sex).lower()

        if self.sex in ['m', '1']:
            self.sex = '1'
        elif self.sex in ['f', '2']:
            self.sex = '2'
        else:
            raise TypeError('define sex as m or f')

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

        row = df2.iloc[
            (df2['Length'] - self.length).abs().argsort()[:1]
        ]

        L = float(row['L'].values[0])
        M = float(row['M'].values[0])
        S = 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

        return {
            "Z": round(z, 2),
            "Percentile": round(percentile, 2)
        }
