#!/usr/bin/env python3
"""Synthetic arithmetic, not a forecasting model or library benchmark. Python3 stdlib."""
import csv
import json
from collections import defaultdict
from fractions import Fraction
from pathlib import Path

HERE = Path(__file__).resolve().parent


def load(path=HERE / 'hierarchy.csv'):
    with path.open(encoding='utf-8', newline='') as stream:
        return list(csv.DictReader(stream))


def reconcile(rows, scenario):
    selected = [dict(row) for row in rows if row['scenario'] == scenario]
    if not selected:
        raise ValueError('unknown scenario')
    if len({(row['period'], row['unit']) for row in selected}) != 1:
        raise ValueError('mixed period or unit')
    seen = set()
    for row in selected:
        key = (row['level'], row['channel'], row['sku'])
        if key in seen:
            raise ValueError('duplicate node / SKU-channel-period')
        seen.add(key)
        if row['level'] not in {'brand', 'channel', 'sku'}:
            raise ValueError('unknown level')
        row['value'] = Fraction(row['forecast'])
        if row['value'] < 0:
            raise ValueError('negative forecast unsupported by this example')
        if row['level'] == 'brand' and (row['channel'] or row['sku']):
            raise ValueError('malformed brand')
        if row['level'] == 'channel' and (not row['channel'] or row['sku']):
            raise ValueError('malformed channel')
        if row['level'] == 'sku' and (not row['channel'] or not row['sku']):
            raise ValueError('malformed leaf')
    roots = [r for r in selected if r['level'] == 'brand']
    channels = {r['channel']: r['value'] for r in selected if r['level'] == 'channel'}
    leaves = {(r['channel'], r['sku']): r['value'] for r in selected if r['level'] == 'sku'}
    if len(roots) != 1 or not channels or not leaves:
        raise ValueError('incomplete hierarchy')
    sums = defaultdict(Fraction)
    for (channel, _sku), value in leaves.items():
        if channel not in channels:
            raise ValueError('leaf without channel parent')
        sums[channel] += value
    if set(sums) != set(channels):
        raise ValueError('channel without leaf')
    if dict(sums) != channels:
        raise ValueError('independent channel differs from its leaves: resolve this separately')
    total = sum(leaves.values(), Fraction(0))
    if total <= 0:
        raise ValueError('zero total cannot define proportional shares')
    top = roots[0]['value']
    topdown = {key: top * value / total for key, value in leaves.items()}
    tdchannels = {channel: sum(v for (c, _), v in topdown.items() if c == channel)
                  for channel in channels}
    return {'top': top, 'bottom_up': total, 'gap': top-total,
            'bottom_up_channels': dict(sums), 'bottom_up_leaves': leaves,
            'top_down_channels': tdchannels, 'top_down_leaves': topdown}


def as_number(value):
    return str(value.numerator) if value.denominator == 1 else str(value)


def write_example(rows):
    with (HERE / 'reconciliation-example.csv').open('w', encoding='utf-8', newline='') as stream:
        writer = csv.writer(stream)
        writer.writerow(['scenario', 'method', 'period', 'unit', 'level', 'channel', 'sku', 'forecast'])
        for scenario in ('base', 'extended'):
            result = reconcile(rows, scenario)
            period, unit = next((r['period'], r['unit']) for r in rows if r['scenario'] == scenario)
            for method in ('bottom_up', 'top_down'):
                total = result['bottom_up'] if method == 'bottom_up' else result['top']
                writer.writerow([scenario, method, period, unit, 'brand', '', '', as_number(total)])
                for channel, value in result[method+'_channels'].items():
                    writer.writerow([scenario, method, period, unit, 'channel', channel, '', as_number(value)])
                for (channel, sku), value in result[method+'_leaves'].items():
                    writer.writerow([scenario, method, period, unit, 'sku', channel, sku, as_number(value)])


if __name__ == '__main__':
    rows = load()
    write_example(rows)
    print(json.dumps({name: {k: as_number(v) for k, v in reconcile(rows, name).items()
                            if isinstance(v, Fraction)} for name in ('base', 'extended')},
                     ensure_ascii=False, indent=2))