"""
doe.py

This module provides functions to support optimized experimental design (doe).

Main functions:
- entropy(sample): Computes entropy from a sample.

"""

####################
## MODULE IMPORTS ##
####################

import numpy as np
import pandas as pd

##############
## FUNCTION ##
##############

## entropy calculation ##
#########################

def entropy_from_hist(sample, num_bins):
    """
    Calculate the entropy of a sample using histogram binning.

    Parameters:
    sample ((n,)-array): Input data sample.
    num_bins (int): Number of bins to use for histogram estimation.

    Returns:
    float: Estimated entropy value.
    """

    # Compute histogram: counts represents the number of samples in each bin,
    # bin_edges gives the edges of the bins.
    counts, bin_edges = np.histogram(sample, bins=num_bins)

    # Compute the bin width (equal bin widths)
    bin_width = bin_edges[1] - bin_edges[0]

    # Normalize counts to obtain a discrete probability distribution
    probabilities = counts / len(sample)

    # Remove zero probabilities to avoid log(0) issues
    probabilities = probabilities[probabilities > 0]

    # Compute entropy using base-2 logarithm (measured in bits)
    entropy_hist = -np.sum(probabilities * np.log2(probabilities))

    # Add log2(bin_width) to account for the continuous nature of the data
    entropy_hist += np.log2(bin_width)

    return entropy_hist  # Return the estimated entropy value

def entropy(sample):
    """
    Compute entropy of a sample using different binning methods.

    Parameters:
    sample ((n,)-array): Input data sample.

    Returns:
    pd.DataFrame: A DataFrame containing the number of bins and entropy 
                  values for different binning methods.
    """
    
    ## Initialize output DataFrame with indices 'bins' and 'H' (entropy values),
    ## and columns for different binning methods.
    H = pd.DataFrame(index=['bins', 'H'], columns=['Freedman', 'Sturges', 'Hacine', 'SQRT'])
    
    ## Calculate the number of bins using different binning methods
    H.loc['bins', 'Freedman'] = num_bins_freedman(sample)  # Freedman-Diaconis rule
    H.loc['bins', 'Sturges'] = num_bins_sturges(sample)    # Sturges' formula
    H.loc['bins', 'Hacine'] = num_bins_hacine(sample)      # Hacine's method
    H.loc['bins', 'SQRT'] = num_bins_sqrt(sample)          # Square-root rule
    
    ## Compute entropy for each binning method using the entropy_from_hist function
    for method in H.columns:
        H.loc['H', method] = entropy_from_hist(sample, H.loc['bins', method])
        
    ## Return the DataFrame containing the number of bins and entropy values
    return H


def entropy_normal(s):
    """
    Compute the entropy of a normal distribution in bits.

    Parameters:
    s (float): Standard deviation of the normal distribution.

    Returns:
    float: Entropy of the normal distribution in bits.
    """
    
    # Compute entropy using the analytical formula (Cover & Thomas 2006, Table 17.1; Michalowicz 2014, Table 2.3)
    H = 1/2 * np.log2(2 * np.pi * np.e * s**2)
    
    return H  # Return the entropy value

import numpy as np  # Import NumPy for mathematical operations

def entropy_lognormal(mln, sln):
    """
    Compute the entropy of a lognormal distribution in bits.

    Parameters:
    mln (float): Mean of the natural logarithm of the lognormal distribution.
    sln (float): Standard deviation of the natural logarithm of the lognormal distribution.

    Returns:
    float: Entropy of the lognormal distribution in bits.
    """
    
    # Compute entropy using the analytical formula (Cover & Thomas 2006, Table 17.1; Kvålseth 1982, Eqn. 8)
    H = (1/2) * np.log2(2 * np.pi * np.e * (sln**2)) + mln * np.log2(np.e)
    
    return H  # Return the entropy value





#########
## AUX ##
#########

def num_bins_freedman(sample):
    """
    Calculate the optimal number of bins for a histogram using the Freedman-Diaconis rule.
    D. Freedman and P. Diaconis, “On the histogram as a density estimator: L2 theory,” Z. Wahrscheinlichkeitstheorie verw Gebiete, vol. 57, no. 4, pp. 453–476, Dec. 1981, doi: 10.1007/BF01025868.

    Input parameters:
    - sample (np.array): Input data sample

    Returns:
    - num_bins (int): Recommended number of bins
   
   # Example usage:
    sample = np.random.normal(0, 1, 1000)  
    num_bins = freedman_bins(sample)
    
    """
    Q1 = np.percentile(sample, 25)
    Q3 = np.percentile(sample, 75)
    IQR = Q3 - Q1
    h = 2 * IQR / np.cbrt(len(sample))  # Bin width
    num_bins_freedman = int(np.ceil((sample.max() - sample.min()) / h))
    return num_bins_freedman 

def num_bins_sturges(sample):
    """
    Calculate the optimal number of bins for a histogram using the Sturges’ rule.
    H. O. Lohaka, “Making a Grouped-data Frequency Table: Development and Examination of the Iteration Algorithm,” Doctoral dissertation, Ohio University, Ohio, USA, 2007. Accessed: Jan. 30, 2025. [Online].

    Input parameters:
    - sample (np.array): Input data sample

    Returns:
    - num_bins (int): Recommended number of bins
   
   # Example usage:
    sample = np.random.normal(0, 1, 1000) 
    num_bins = sturges_bins(sample)
    
    """
    num_bins_sturges = int(np.ceil(np.log2(len(sample)) + 1))
    return num_bins_sturges

def num_bins_hacine(sample):
    """
    Calculate the optimal number of bins for a histogram using the Abdenour Hacine-Gharbi et al. (2012) rule:
    A. Hacine-Gharbi, P. Ravier, R. Harba, and T. Mohamadi, “Low bias histogram-based estimation of mutual information for feature selection,” Pattern Recognition Letters, vol. 33, no. 10, pp. 1302–1308, Jul. 2012, doi: 10.1016/j.patrec.2012.02.022.

    Input parameters:
    - sample (np.array): Input data sample

    Returns:
    - num_bins (int): Recommended number of bins
   
   # Example usage:
    sample = np.random.normal(0, 1, 1000) 
    num_bins = hacine_bins(sample)
    
    """
    no_of_sample = len(sample)
    xi = (8 + 324 * no_of_sample + 12 * np.sqrt(36 * no_of_sample + 729 * no_of_sample**2))**(1/3)
    hacine_bins = int(round((xi / 6) + (2 / (3 * xi)) + (1 / 3), 0))
    return hacine_bins

def num_bins_sqrt(sample):
    """
    Calculate the optimal number of bins for a histogram using the square-root rule.
    H. O. Lohaka, “Making a Grouped-data Frequency Table: Development and Examination of the Iteration Algorithm,” Doctoral dissertation, Ohio University, Ohio, USA, 2007. Accessed: Jan. 30, 2025. [Online].

    Input parameters:
    - sample (np.array): Input data sample

    Returns:
    - num_bins (int): Recommended number of bins
   
   # Example usage:
    sample = np.random.normal(0, 1, 1000) 
    num_bins = sqrt_bins(sample)
    
    """
    sqrt_bins = int(np.sqrt(len(sample)))
    return sqrt_bins