"""Offline account-cluster planning approximation for E1 activation. Standard library only.

Decision rule (sized here): roll out only if the lower bound of the two-sided (1 - alpha) interval for the
activation gain is above zero. The practical threshold is the true gain the trial is sized to detect.
"""
import argparse
import json
import math
from pathlib import Path
from statistics import NormalDist

NORMAL = NormalDist()


def _finite(*values):
    return all(not isinstance(v, bool) and isinstance(v, (int, float)) and math.isfinite(v) for v in values)


def _positive_int(value, label):
    if isinstance(value, bool) or not isinstance(value, int) or value < 1:
        raise ValueError(label + ' must be a positive integer')


def design_effect(users_per_account, icc, size_cv=0.0):
    """Variance inflation of the user-weighted mean: 1 + ((1 + CV^2) * mean size - 1) * ICC.

    Exact for a fixed-weight user average over independent accounts with exchangeable correlation, where CV is the
    population coefficient of variation of account sizes. With CV 0 it is the equal-cluster 1 + (m - 1) * ICC.
    """
    if not _finite(icc, size_cv) or not 0 <= icc <= 1 or size_cv < 0:
        raise ValueError('Invalid correlation or cluster-size variation')
    _positive_int(users_per_account, 'Account size')
    return 1 + ((1 + size_cv ** 2) * users_per_account - 1) * icc


def _variances(p0, p1):
    pooled = (p0 + p1) / 2
    return 2 * pooled * (1 - pooled), p0 * (1 - p0) + p1 * (1 - p1)


def plan(p0, p1, alpha, power, users_per_account, icc, size_cv=0.0):
    """Accounts per arm for the lower-bound rule to have `power` at a true gain of p1 - p0."""
    if not _finite(p0, p1, alpha, power, icc):
        raise ValueError('Parameters must be finite numbers')
    if not (0 < p0 < 1 and 0 < p1 < 1 and p0 != p1 and 0 < alpha < .5 and .5 < power < 1 and 0 <= icc <= 1):
        raise ValueError('Invalid rate, effect, error probability, power or correlation')
    effect = design_effect(users_per_account, icc, size_cv)
    null_variance, alternative_variance = _variances(p0, p1)
    z_alpha = NORMAL.inv_cdf(1 - alpha / 2)
    z_power = NORMAL.inv_cdf(power)
    independent = (z_alpha * math.sqrt(null_variance) + z_power * math.sqrt(alternative_variance)) ** 2 / (p1 - p0) ** 2
    accounts = math.ceil(independent * effect / users_per_account)
    return {'p0': p0, 'p1': p1, 'alpha': alpha, 'target_power': power,
            'users_per_account': users_per_account, 'icc': icc, 'size_cv': size_cv,
            'independent_users_per_arm_unrounded': independent,
            'design_effect': effect, 'accounts_per_arm': accounts,
            'total_accounts': accounts * 2, 'total_users': accounts * users_per_account * 2,
            'null_variance_sum': null_variance, 'alternative_variance_sum': alternative_variance}


def rollout_probability(p0, gain, alpha, accounts_per_arm, users_per_account, icc, size_cv=0.0, estimate_threshold=None):
    """Chance of a rollout decision at a true gain, for a given design.

    Lower-bound rule: the estimate exceeds z(1 - alpha/2) times its pooled-variance standard error.
    With estimate_threshold set, the estimate must also reach that threshold (the earlier E1 rule).
    Normal approximation; a significant result in the wrong direction is not a rollout.
    """
    p1 = p0 + gain
    if not _finite(p0, gain, alpha) or not (0 < p0 < 1 and 0 < p1 < 1 and gain >= 0 and 0 < alpha < .5):
        raise ValueError('Invalid rate, gain or error probability')
    _positive_int(accounts_per_arm, 'Accounts per arm')
    inflation = design_effect(users_per_account, icc, size_cv) / (accounts_per_arm * users_per_account)
    null_variance, alternative_variance = _variances(p0, p1)
    bar = NORMAL.inv_cdf(1 - alpha / 2) * math.sqrt(null_variance * inflation)
    if estimate_threshold is not None:
        bar = max(bar, estimate_threshold)
    return 1 - NORMAL.cdf((bar - gain) / math.sqrt(alternative_variance * inflation))


def accounts_for_estimate_rule(p0, gain, threshold, alpha, power, users_per_account, icc):
    """Accounts per arm for the earlier rule (interval above zero AND estimate at least `threshold`) to reach `power`."""
    if not _finite(gain, threshold) or not 0 < threshold < gain:
        raise ValueError('The true gain must exceed the estimate threshold')
    base = plan(p0, p0 + gain, alpha, power, users_per_account, icc)
    z_power = NORMAL.inv_cdf(power)
    threshold_users = (z_power * math.sqrt(base['alternative_variance_sum'])) ** 2 / (gain - threshold) ** 2
    independent = max(base['independent_users_per_arm_unrounded'], threshold_users)
    accounts = math.ceil(independent * base['design_effect'] / users_per_account)
    binding = 'estimate threshold' if threshold_users > base['independent_users_per_arm_unrounded'] else 'interval above zero'
    return {'true_gain': gain, 'estimate_threshold': threshold, 'independent_users_per_arm_unrounded': independent,
            'binding_condition': binding, 'accounts_per_arm': accounts, 'total_accounts': accounts * 2,
            'total_users': accounts * users_per_account * 2}


def calculate(assumptions):
    if assumptions['allocation'] != '1:1 accounts' or assumptions['attrition'] != 0:
        raise ValueError('Only equal allocation and zero attrition are implemented')
    if assumptions.get('decision_rule', 'lower_bound_above_zero') != 'lower_bound_above_zero':
        raise ValueError('Only the lower-bound rule is sized')
    hypotheses = assumptions['primary_hypotheses']
    _positive_int(hypotheses, 'Primary hypothesis count')
    if 'activation_completion_days' in assumptions:
        for key in ('activation_completion_days', 'reopen_observation_days', 'followup_days'):
            _positive_int(assumptions[key], key)
        if assumptions['activation_completion_days'] + assumptions['reopen_observation_days'] != assumptions['followup_days']:
            raise ValueError('Activation window and reopen observation must add up to the follow-up')
    p0, p1, m, power = assumptions['p0'], assumptions['p1'], assumptions['users_per_account'], assumptions['power']
    threshold = assumptions.get('practical_threshold', p1 - p0)
    if not _finite(threshold) or not math.isclose(threshold, p1 - p0, abs_tol=1e-12):
        raise ValueError('The sized gain must equal the practical threshold')
    alpha = assumptions['family_alpha'] / hypotheses
    rows = [plan(p0, p1, alpha, power, m, icc) for icc in assumptions['icc_values']]
    if len({row['icc'] for row in rows}) != len(rows) or assumptions['base_icc'] not in [r['icc'] for r in rows]:
        raise ValueError('Sensitivity rows must be unique and include the base ICC')
    result = {'status': 'Hypothetical planning calculation, not observed outcomes',
              'method': 'Two-sided normal approximation; exchangeable account clusters',
              'base_icc': assumptions['base_icc'], 'rows': rows}
    if 'size_cv_values' not in assumptions:
        return result
    base = next(r for r in rows if r['icc'] == assumptions['base_icc'])
    cvs = assumptions['size_cv_values']
    if len(set(cvs)) != len(cvs) or any(not _finite(cv) or cv <= 0 for cv in cvs):
        raise ValueError('Cluster-size CV rows must be unique and positive')
    gains = assumptions['operating_characteristic_gains']
    if len(set(gains)) != len(gains) or 0 not in gains or not any(math.isclose(g, threshold) for g in gains):
        raise ValueError('Operating characteristics must include zero and the practical threshold')
    result['decision_rule'] = 'Roll out only if the lower bound of the two-sided (1-alpha) interval is above zero'
    result['practical_threshold'] = threshold
    result['cluster_size_rows'] = [plan(p0, p1, alpha, power, m, assumptions['base_icc'], cv) for cv in cvs]
    result['operating_characteristics'] = {
        'accounts_per_arm': base['accounts_per_arm'], 'users_per_account': m, 'icc': base['icc'],
        'rows': [{'true_gain': g, 'p1': p0 + g,
                  'lower_bound_rule': rollout_probability(p0, g, alpha, base['accounts_per_arm'], m, base['icc']),
                  'earlier_estimate_threshold_rule': rollout_probability(
                      p0, g, alpha, base['accounts_per_arm'], m, base['icc'], estimate_threshold=threshold)}
                 for g in gains]}
    result['estimate_threshold_alternative'] = [
        accounts_for_estimate_rule(p0, g, threshold, alpha, power, m, base['icc'])
        for g in assumptions['estimate_rule_planning_gains']]
    return result


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--assumptions', type=Path, default=Path(__file__).with_name('assumptions.json'))
    args = parser.parse_args()
    print(json.dumps(calculate(json.loads(args.assumptions.read_text(encoding='utf-8'))), indent=2))


if __name__ == '__main__':
    main()