#!/usr/bin/env python
## Axiomatic Design Information Calculator (and plotter)
## Author: Joseph Timothy foley <foley AT RU.IS>
## Start Date:  2026-02-27
## Input:  data in csv file
## Output:  information calculation and PDF for report/presentation
import os
import logging
import argparse
from pathlib import PurePath##https://docs.python.org/3/library/pathlib.html#module-pathlib
import numpy as np
import matplotlib
import matplotlib.pyplot as plt
from scipy.stats import norm,t
import scipy.stats
import pandas as pd

#Main program loop
print("""Axiomatic Design Information Calculator by Joseph. T. Foley<foley AT ru DOT is>
From https://gitea.cs.ru.is/AxiomaticDesign/adcalc/""")
parser = argparse.ArgumentParser(
    description="Axiomatic Design Information Calculator.")
subparsers = parser.add_subparsers(dest='mode')
subparsers.required = True

### MODE DATA
parser_data = subparsers.add_parser("DATA")
parser_data.add_argument('csvfile',
                    help="CSV file with data and headers")
parser_data.add_argument('column',
                    help='Which column header to take data from')
## MODE SIM
parser_sim = subparsers.add_parser("SIM")
parser_sim.add_argument('samplesize', type=int,
                    help="sample size")
parser_sim.add_argument('mean', type=float,
                    help="mean(average) value")
parser_sim.add_argument('stddev', type=float,
                    help="sample standard deviation")
## General Arguments
parser.add_argument('--lowerbound', type=float,
                    help='Tolerance low limit')
parser.add_argument('--upperbound', type=float,
                    help='Tolerance high limit')
parser.add_argument('--normalizey', action="store_true",
                    help='Set y-axis to normalized probability density')
parser.add_argument('--log', default="INFO",
                    help='Console log level:  Number or DEBUG, INFO, WARNING, ERROR')
parser.add_argument('--legend', action="store_true",
                    help='Put legend on the PDF graph')
parser.add_argument('--graphinfo', action="store_true",
                    help='Put information on the PDF graph')
parser.add_argument('--xlabel',
                    help='X-axis label, if needed')
parser.add_argument('--outfile',
                    help="output graph to PDF file")
parser.add_argument('--fontsize', default=14, type=int,
                    help="Adjust font size")

args = parser.parse_args()

## Set up logging
numeric_level = getattr(logging, args.log.upper(), None)
if not isinstance(numeric_level, int):
    raise ValueError(f'Invalid log level: {args.log}')
#print(f"Log level:  {numeric_level}")
logger = logging.getLogger("app")
logger.setLevel(numeric_level)
# log everything to file
logpath = "infocalc.log"
fh = logging.FileHandler(logpath)
fh.setLevel(logging.DEBUG)
# log to console
ch = logging.StreamHandler()
ch.setLevel(numeric_level)
# create formatter and add to handlers
consoleformatter = logging.Formatter('%(message)s')
ch.setFormatter(consoleformatter)
spamformatter = logging.Formatter('%(asctime)s %(name)s[%(levelname)s] %(message)s')
fh.setFormatter(spamformatter)
# add the handlers to logger
logger.addHandler(ch)
logger.addHandler(fh)

logger.debug("Creating infocalc log file %s", logpath)

# seed values for variable scoping
mean = 0
stddev = 1
samplesize =1

if args.mode == "DATA":
    # filename pre-processing for output
    inpath = PurePath(args.csvfile)
    print(f"Input: {inpath}")    
    # grab the data and process
    data = np.array(pd.read_csv(inpath)[args.column])
    mean = data.mean()
    stddev = data.std(ddof=1)
    # Delta Degrees of Freedom: ddof=0 for population, ddof=1 for sample std dev
    samplesize = len(data)
elif args.mode == "SIM":
    mean = args.mean
    stddev = args.stddev
    samplesize = args.samplesize
df = samplesize - 1 
    
prob_upper = 1
prob_lower = 0
if args.upperbound is not None:
    prob_upper = t.cdf(df=df,x=args.upperbound, loc=mean, scale=stddev)
if args.lowerbound is not None:
    prob_lower = t.cdf(df=df,x=args.lowerbound, loc=mean, scale=stddev)
prob = prob_upper - prob_lower

info = -np.emath.log2(prob)
logger.info(f"n={samplesize},df={df},mean={mean},stddev={stddev}")
logger.info(f"P()={prob},I={info} bits")

## set default fontsize
matplotlib.rcParams['font.size']=args.fontsize
## place text on plot:  https://matplotlib.org/3.3.4/gallery/recipes/placing_text_boxes.html
fig, ax = plt.subplots()
if args.graphinfo:#put info on corner of graph
    textstr = '\n'.join((
        r'$n=%d$' % (samplesize),
        r'$\mu=%.2f$' % (mean, ),
        r'$\sigma=%.2f$' % (stddev, ),
        r'$P=%.2f$' % (prob, ),
        r'$I=%.2f$ bits' % (info, )))
    # these are matplotlib.patch.Patch properties
    props = dict(boxstyle='round', facecolor='wheat', alpha=0.5)

    # place a text box in upper left in axes coords
    ax.text(0.05, 0.95, textstr, transform=ax.transAxes, fontsize=args.fontsize,
            verticalalignment='top', bbox=props)
xgraphlimits = {"min": mean-3*stddev, "max": mean+3*stddev}
if args.lowerbound and xgraphlimits["min"] > args.lowerbound:
    xgraphlimits["min"] = args.lowerbound
if args.upperbound and xgraphlimits["max"] < args.upperbound:
    xgraphlimits["max"] = args.upperbound

x = np.linspace(xgraphlimits["min"], xgraphlimits["max"], 500)
y = norm.pdf(x, loc=mean, scale=stddev)
if args.normalizey:
    y = y * stddev#rescale back to unity area
plt.axvline(x=mean, color="green", linestyle="dashed", label="mean")
if args.lowerbound is not None:
    plt.axvline(args.lowerbound, color="red")
if args.upperbound is not None:
    plt.axvline(args.upperbound, color="red")

plt.plot(x, y, 'b-', label='T distribution')

# Filter for which region to fill
coloredregion = x#default fill all
if args.lowerbound is not None and args.upperbound is not None:
    coloredregion = (x >= args.lowerbound) & ( x <= args.upperbound )
elif args.upperbound is not None:
    coloredregion =  x <= args.upperbound 
elif args.lowerbound is not None:
    coloredregion = x >= args.lowerbound
    
plt.fill_between(x, 0, y, where=coloredregion, color="grey", alpha=0.5, label="Design range",)
    
if args.xlabel:
    plt.xlabel(args.xlabel)
plt.ylabel('Probability density')
if args.legend:
    plt.legend()
#plt.grid(True)
top = plt.ylim()[1]

if args.outfile:
    logger.info(f"Graph output to {args.outfile}")
    plt.savefig(args.outfile,bbox_inches='tight')
else:
    plt.show()
