#!/usr/bin/env python

"""Distribution functions

This module provides functions for dealing with normal distributions
and generating error maps.

When called directly as main, it allows for converting a threshold map
into an error map.

```
$ python -m mlcsim.dist --help

usage: dist.py [-h] [-b {1,2,3,4}] -f F [-o O]

options:
  -h, --help    show this help message and exit
  -b {1,2,3,4}  bits per cell
  -f F          Threshold map json to convert
  -o O          output to file
```
"""

import argparse
import json
from pprint import pprint
from typing import Dict, List

import numpy as np
from scipy import stats as ss  # type: ignore

# https://stackoverflow.com/a/32574638/9047818
# https://stackoverflow.com/a/13072714/9047818
def normalMidpoint(mean_a: float, mean_b: float, std_a: float, std_b: float) -> float:
    """Find the midpoint between two normal distributions

    Args:
        mean_a (float): Mean of first distribution
        mean_b (float): Mean of second distribution
        std_a (float): Std dev of first distribution
        std_b (float): Std dev of second distribution

    Returns:
        float: Midpoint between distributions
    """
    a = 1 / (2 * std_a**2) - 1 / (2 * std_b**2)
    b = mean_b / (std_b**2) - mean_a / (std_a**2)
    c = (
        mean_a**2 / (2 * std_a**2)
        - mean_b**2 / (2 * std_b**2)
        - np.log(std_b / std_a)
    )
    roots = np.roots([a, b, c])
    masked = np.ma.masked_outside(roots, mean_a, mean_b)
    return float(masked[~masked.mask][0][0])


# https://www.askpython.com/python/normal-distribution
def normalChance(mean: float, stdev: float, thr: float) -> float:
    """Find the chance of a normal distribution above/below a given value

    Args:
        mean (float): Mean of the distribution
        stdev (float): Std dev of the distribution
        thr (float): Threshold to check above/below

    Returns:
        float: Chance for threshold to end up above/below the given point in the distribution
    """
    chance = ss.norm(loc=mean, scale=stdev).cdf(thr)
    return float(chance if mean > thr else 1 - chance)


def genErrorMap(thr_maps: Dict[str, List[List[float]]], bpc: int) -> List[List[float]]:
    """Generate an error map from a threshold map

    Args:
        thr_maps (dict): Threshold map
        bpc (int): Bits per cell

    Raises:
        ValueError: if the given bpc is not in the threshold map

    Returns:
        list: Error map from the threshold map
    """
    if str(bpc) not in thr_maps.keys():
        raise ValueError(f"Threshold map does not have values for {bpc} levels")
    thr_map: List[List[float]] = thr_maps[str(bpc)]
    err_map = [[0.0]]

    for i in range(len(thr_map) - 1):
        mid = normalMidpoint(
            thr_map[i][0], thr_map[i + 1][0], thr_map[i][1], thr_map[i + 1][1]
        )
        up = normalChance(thr_map[i][0], thr_map[i][1], mid)
        dn = normalChance(thr_map[i + 1][0], thr_map[i + 1][1], mid)

        err_map[i].append(up)
        err_map.append([dn])

    err_map[-1].append(0.0)

    return err_map


def _main():
    parser = argparse.ArgumentParser()

    parser.add_argument(
        "-b", type=int, default=2, choices=range(1, 5), help="bits per cell"
    )
    parser.add_argument("-f", required=True, help="Threshold map json to convert")
    parser.add_argument("-o", type=str, help="output to file")

    args = parser.parse_args()

    with open(args.f) as f:
        thr_map = json.load(f)

    err_map = genErrorMap(thr_map, args.b)

    if args.o:
        with open(args.o, "w") as f:
            json.dump(err_map, f)
    else:
        pprint(err_map)


if __name__ == "__main__":
    _main()

[((53, 12, 53, 31), 'numpy.roots', 'np.roots', ({(53, 21, 53, 30): '[a, b, c]'}, {}), '([a, b, c])', True, 'import numpy as np\n'), ((54, 13, 54, 56), 'numpy.ma.masked_outside', 'np.ma.masked_outside', ({(54, 34, 54, 39): 'roots', (54, 41, 54, 47): 'mean_a', (54, 49, 54, 55): 'mean_b'}, {}), '(roots, mean_a, mean_b)', True, 'import numpy as np\n'), ((108, 13, 108, 38), 'argparse.ArgumentParser', 'argparse.ArgumentParser', ({}, {}), '()', False, 'import argparse\n'), ((51, 10, 51, 31), 'numpy.log', 'np.log', ({(51, 17, 51, 30): '(std_b / std_a)'}, {}), '(std_b / std_a)', True, 'import numpy as np\n'), ((119, 18, 119, 30), 'json.load', 'json.load', ({(119, 28, 119, 29): 'f'}, {}), '(f)', False, 'import json\n'), ((127, 8, 127, 23), 'pprint.pprint', 'pprint', ({(127, 15, 127, 22): 'err_map'}, {}), '(err_map)', False, 'from pprint import pprint\n'), ((70, 13, 70, 43), 'scipy.stats.norm', 'ss.norm', (), '', True, 'from scipy import stats as ss\n'), ((125, 12, 125, 33), 'json.dump', 'json.dump', ({(125, 22, 125, 29): 'err_map', (125, 31, 125, 32): 'f'}, {}), '(err_map, f)', False, 'import json\n')]