# coding: utf-8
"""A monte carlo dice simulator

Given a method for rolling a set of dice, generate a histogram of results by
actually trying it a bunch.
"""
import argparse
import importlib
import os
import sys
from collections import defaultdict

from dice import d4, d6, d8, d10, d12, d20, d100, D  # noqa: F401


def mk_roller(*dice, drop=0):
    """Create a roller from a set of dice and a number to drop

    If drop >= 0, drop the lowest rolls. If drop < 0, drop the highest.
    """
    def roller():
        rolls = [d.roll() for d in dice]
        keep = len(dice) - abs(drop)
        rolls = list(sorted(rolls, reverse=drop >= 0))[0:keep]
        return sum(rolls)
    roller.__name__ = [str(d) for d in dice].__str__() + ' drop %d' % drop
    return roller


adv = mk_roller(d20, d20, drop=1)
adv.__doc__ = 'advantage'
disadv = mk_roller(d20, d20, drop=-1)
disadv.__doc__ = 'disadvantage'


def noam():
    """4d6 + 1d4 drop 2"""
    rolls = [d6.roll(), d6.roll(), d6.roll(), d6.roll(), d4.roll()]
    rolls = list(sorted(rolls, reverse=True))
    rolls = rolls[0:3]
    return sum(rolls)


def threed6():
    """3d6"""
    return 3 * d6


def fourdrop1():
    """4d6 drop 1"""
    rolls = sorted([d6.roll() for _ in range(4)])
    top3 = list(reversed(rolls))[0:3]
    return sum(top3)


def oned20():
    """1d20 to embrace chaos"""
    return 1 * d20


def r4d6ro1dl1():
    """4d6, reroll 1s once, drop lowest"""
    rolls = []
    for _ in range(4):  # 4d6
        roll = d6.roll()
        if roll == 1:  # Reroll 1s once
            roll = d6.roll()
        rolls.append(roll)
    rolls = list(sorted(rolls, reverse=True))
    return sum(rolls[0:3])


TESTS = {
    '3d6': threed6,
    '4d6d1': fourdrop1,
    'adv': adv,
    'd20': oned20,
    'disadv': disadv,
    'noam': noam,
    '4d6ro1d1': r4d6ro1dl1,
}


def calc_stats(totals, iters):
    roll_total = sum(k * v for k, v in totals.items())
    mean = roll_total / iters

    variance = 0.0
    for score, count in totals.items():
        variance += ((score - mean) ** 2) * count
    variance /= (iters - 1)
    stdev = variance ** 0.5

    return {
        'mean': mean,
        'stdev': stdev,
        'variance': variance,
    }


def rolltest(roller, iters=10000, precision=3, graph='normal', zoom=1,
             max_cols=None):
    print('Method: %s\nDesc:  %s' % (roller.__name__, roller.__doc__))
    print('Iterations: %d, Precision: %d' % (iters, precision))

    totals = defaultdict(int)
    for _ in range(iters):
        score = roller()
        totals[score] += 1

    stats = calc_stats(totals, iters)
    print('Average: {mean:0.2f}; StDev: {stdev:0.2f}'.format(**stats))

    output_fmt = '{0:3d} [{1:5.1f}%]: {2}'
    if max_cols is None:
        prefix_length = len(output_fmt.format(18, 15.1, ''))
        max_cols = os.get_terminal_size().columns - prefix_length

    if graph == 'normal':
        for k in sorted(totals.keys()):
            frac = round(totals[k] / iters, precision)
            cols = int(frac * max_cols * zoom)
            print(output_fmt.format(k, frac * 100., '*' * cols))
    elif graph == 'atmost':
        cumulative = 0.0
        for k in sorted(totals.keys()):
            cumulative += round(totals[k] / iters, precision)
            cols = int(cumulative * max_cols)
            print(output_fmt.format(k, cumulative * 100., '*' * cols))
    elif graph == 'atleast':
        cumulative = 1.0
        for k in sorted(totals.keys()):
            cols = int(cumulative * max_cols)
            print(output_fmt.format(k, cumulative * 100., '*' * cols))
            cumulative -= round(totals[k] / iters, precision)
    else:
        print('Unsupported graph mode: %s' % graph)


def make_parser():
    parser = argparse.ArgumentParser()
    parser.add_argument('-p', '--precision', default=3, type=int)
    parser.add_argument('-i', '--iters', default=10000, type=int)
    parser.add_argument('-t', '--test', dest='tests', choices=TESTS.keys(),
                        action='append', default=[])
    parser.add_argument('-l', '--lambda', dest='lamb')
    parser.add_argument('-r', '--raw')
    parser.add_argument('--imp', help=('A dotted path to a roller, e.g. '
                                       '--imp my.module.roller, where '
                                       'roller is a callable'))
    parser.add_argument('-g', '--graph', default='normal',
                        choices=['normal', 'atleast', 'atmost'])
    parser.add_argument('-z', '--zoom', default=1, type=int)
    parser.add_argument('-m', '--max-cols', default=None, type=int)
    return parser


def get_opts(argv=None):
    parser = make_parser()
    if argv is None:
        argv = sys.argv[1:]
    return parser.parse_args(argv)


def main():
    options = get_opts()

    to_run = []
    for test in options.tests:
        to_run.append(TESTS[test])

    if options.lamb:
        roller = eval('lambda: %s' % options.lamb)
        roller.__doc__ = 'user lambda: %s' % options.lamb
        to_run.append(roller)

    if options.raw:
        roller = eval(options.raw)
        roller.__doc__ = 'user input: %s' % options.raw
        to_run.append(roller)

    if options.imp:
        path, _, method = options.imp.rpartition('.')
        module = importlib.import_module(path)
        roller = getattr(module, method)
        to_run.append(roller)

    for test in to_run:
        rolltest(test, iters=options.iters, precision=options.precision,
                 graph=options.graph, zoom=options.zoom,
                 max_cols=options.max_cols)


if __name__ == '__main__':
    main()