#  Python 3
#  File: Assignment1.py
#  Demo Python code for Computational Statistics class, Fall 2026
#  Authors: Jonathan Goodman, goodman@cims.nyu.edu
#  class web site: https://math.nyu.edu/~goodman/teaching/ComputationalStatistics2026/ComputationalStatistics.html

"""  
    Run n simulated experiments to generate fake data.  
    Use a histogram to estimate the PDF of the outcome
    and the corresponding potential.
    Plot the PDF and potential estiamtes on the same graph and save the 
    plot.
"""

import numpy             as np    # load the nympy library, call it np
import matplotlib.pyplot as plt   # the plotting package

print("Histogram plotting demo.")

#  Set parameters

n       = 500000   # number of experiments (samples)
h_min   = 5        # minimum number of hits in a bin to estimate log(Pr)
x_max   = 6.       # discard samples X_k > x_max, bin simples with X_k < x_max
x_min   = - x_max  # discard samples X_k < x_min.  Chosen symmetrically for now
n_bins  = 50       # number of histogram bins
bin_min = 5        # minimum bin count for potential estimate

#   Calculated parameters, initialize the histogram

dx = (x_max - x_min )/n_bins   # bin size.  Equal sized bins spanning the binning range
bin_centers = np.linspace(x_min + dx/2, x_max - dx/2, n_bins)
bin_counts  = np.zeros( n_bins, dtype = np.int32)

rng = np.random.default_rng(18)    # create an random number generator instance

#  Replace this with the more complex experiment that returns the sample variance from M samples
#  The parameter M will be part of the "params" dictionary data structure defined below.

def experiment( rng, params):
    """create and return the result of one random experiment
       inputs: 
          rng: a random number generator instance
          params: a dictionary with parameters for the 
                  computational experiment
       outputs: 
                  one number, the result of one experiment
    """
    return rng.normal( loc= params["mu"], scale=params["sig"])

#  Do the n experiments and record the results in teh samples array

samples = np.zeros(n)           # will be the results of the n experiments
params = { "mu":0., "sig":1.}   # store the parameters for the run,
for i in range(n):
    samples[i] = experiment(rng,params)    # do the experiments one by one, for now

#   Create bin counts for the samples array

for sample in samples:
           #   only bin samples inside the histogram range
    if ( sample > x_min ) and ( sample < x_max ):
       bin = int( ( sample - x_min )/dx)  # find the index of the bin for this sample
       bin_counts[bin] += 1               # record a hit in that bin

PDF = bin_counts/(n*dx)    # estimated PDF, see assignment1.pdf

pot     = np.zeros(n_bins)  # will hold estimates of - log(PDF), the potential
pot_pts = np.zeros(n_bins)  # will hold the corresponding bin centers

pot_indx = 0                # index into the pot and pot_pts arrays
for i in range(n_bins):
        # estimate the potential for a bin only if there enough hits in that bin
   if bin_counts[i] >= bin_min:
      pot[pot_indx]     = - np.log( PDF[i])
      pot_pts[pot_indx] = bin_centers[i]
      pot_indx         += 1
pot_vals = pot_indx - 1   # the number of potential estimates made

# put two plots in the same figure.
# first, the estimated PDF, see assignment for details

fig, ax1 = plt.subplots(figsize=(8, 5))
ax1.plot(bin_centers, PDF, color = "black", label = "PDF")
title = "PDF and potential estimate with {n:8.2e} samples and {n_bins:4d} bins"
title = title.format(n=n, n_bins=n_bins)
ax1.set( title = title,
       xlabel = "x",
       ylabel = "PDF (est)")
plt.grid()

# second, the estimated potential function

ax2 = ax1.twinx()  # instantiate a second Axes that shares the same x-axis
ax2.plot( pot_pts[0:pot_indx], pot[0:pot_indx], ".", color = "red", label = "potential")
ax2.set( xlabel = "x",
         ylabel = "potential (est)")
fig.legend(bbox_to_anchor=(.9, .3))
fig.tight_layout()  # otherwise the right y-label is slightly clipped
plt.savefig("Assignment1_density_plot_demo.pdf")
plt.show()