"""Command line.

    python -m feedloop run --policy ucb --rounds 30 --seed 1          one policy -> out/run_ucb.json + out/run_ucb.png
    python -m feedloop compare --rounds 30 --seed 1                   every policy, the mid-run switch and the new item
                                                                      -> out/compare.json + the figures the write-up uses
"""
from __future__ import annotations

import argparse
import json
import os
import sys
import time

import numpy as np

from . import policies
from .sim import METRICS, run
from .world import TOPICS, make_world

OUT = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "out")
PLOTTED = [("diversity", "Topics per feed (of 8)"), ("coverage", "Share of catalogue shown this round"),
           ("gini", "Gini of item exposure"), ("utility", "Feed value vs the best possible"),
           ("bubble", "Share of feed in the user's top topic"), ("clicks", "Clicks per user")]


def add_common(p):
    p.add_argument("--rounds", type=int, default=30)
    p.add_argument("--seed", type=int, default=1)
    p.add_argument("--users", type=int, default=400)
    p.add_argument("--items", type=int, default=300)
    p.add_argument("--k", type=int, default=10)
    p.add_argument("--eps", type=float, default=policies.DEFAULTS["eps"])
    p.add_argument("--ucb-c", type=float, default=policies.DEFAULTS["ucb_c"])
    p.add_argument("--ts-c", type=float, default=policies.DEFAULTS["ts_c"])
    p.add_argument("--mmr-lambda", type=float, default=policies.DEFAULTS["mmr_lambda"])
    p.add_argument("--factors", type=int, default=6)
    p.add_argument("--reg", type=float, default=1.0)
    p.add_argument("--sweeps", type=int, default=2)
    p.add_argument("--repeat-decay", type=float, default=0.7)
    p.add_argument("--no-discount", action="store_true", help="the recommender does not discount items it already showed a user")
    p.add_argument("--new-item-round", type=int, default=0, help="release a held-back item at this round (0 = never)")
    p.add_argument("--switch-round", type=int, default=0)
    p.add_argument("--switch-policy", choices=policies.POLICIES, default=None)
    p.add_argument("--visitor", default=None, help='the browser page\'s ratings as JSON, e.g. \'{"12":1,"40":0}\' (item id: liked), to reproduce a run from the web page exactly')
    p.add_argument("--out", default=OUT)


def settings(a):
    return dict(params={"eps": a.eps, "ucb_c": a.ucb_c, "ts_c": a.ts_c, "mmr_lambda": a.mmr_lambda},
                factors=a.factors, reg=a.reg, sweeps=a.sweeps, repeat_decay=a.repeat_decay,
                system_discount=not a.no_discount)


def visitor_of(a):
    if not a.visitor:
        return None
    return {int(k): bool(int(v)) for k, v in json.loads(a.visitor).items()}


def world_for(a, new_item):
    return make_world(a.seed, a.items, a.users, a.k, new_item=new_item, visitor_ratings=visitor_of(a))


def plot_metrics(runs, path, title=None):
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    fig, axes = plt.subplots(2, 3, figsize=(13, 7.2))
    for ax, (m, label) in zip(axes.ravel(), PLOTTED):
        for name, r in runs:
            ys = r["metrics"][m]
            ax.plot(range(1, len(ys) + 1), ys, label=name, lw=1.8)
            if r.get("switch_round"):
                ax.axvline(r["switch_round"], color="0.6", ls="--", lw=1)
        ax.set_title(label, fontsize=11)
        ax.set_xlabel("round")
        ax.grid(alpha=0.3)
    axes[0, 0].legend(fontsize=9)
    if title:
        fig.suptitle(title)
    fig.tight_layout()
    fig.savefig(path, dpi=130)
    plt.close(fig)


def plot_feed_strip(world, runs, user, path):
    """One user's feed, round by round, coloured by topic, under each policy."""
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    from matplotlib.colors import ListedColormap
    cols = ["#2a78d6", "#eb6834", "#1baf7a", "#eda100", "#e87ba4", "#6250d6", "#8a6d3b", "#008300"]
    fig, axes = plt.subplots(1, len(runs), figsize=(3.2 * len(runs), 4.2), sharey=True)
    for ax, (name, grid) in zip(np.atleast_1d(axes), runs):
        ax.imshow(np.array(grid).T, cmap=ListedColormap(cols), vmin=0, vmax=7, aspect="auto", interpolation="nearest")
        ax.set_title(name, fontsize=11)
        ax.set_xlabel("round")
        ax.set_xticks([0, 9, 19, 29])
        ax.set_xticklabels([1, 10, 20, 30])
    np.atleast_1d(axes)[0].set_ylabel("slot in feed")
    handles = [plt.Rectangle((0, 0), 1, 1, color=cols[t]) for t in range(8)]
    fig.legend(handles, [t["name"] for t in TOPICS], loc="lower center", ncol=8, fontsize=9, frameon=False)
    fig.suptitle(f"User {user}'s feed (top topic: {TOPICS[int(world.top_topic[user])]['name']})", fontsize=12)
    fig.tight_layout(rect=(0, 0.07, 1, 1))
    fig.savefig(path, dpi=130)
    plt.close(fig)


def cmd_run(a):
    os.makedirs(a.out, exist_ok=True)
    w = world_for(a, a.new_item_round > 0)
    t0 = time.time()
    r = run(w, a.policy, a.rounds, switch_round=a.switch_round, switch_policy=a.switch_policy,
            new_item_round=a.new_item_round, visitor_ratings=visitor_of(a), **settings(a))
    r["seconds"] = round(time.time() - t0, 2)
    r["seed"], r["users"], r["items"], r["k"] = a.seed, a.users, a.items, a.k
    base = os.path.join(a.out, f"run_{a.policy}")
    with open(base + ".json", "w") as f:
        json.dump(r, f, indent=1)
    plot_metrics([(policies.LABELS[a.policy], r)], base + ".png")
    m = r["metrics"]
    print(f"{policies.LABELS[a.policy]}: {a.rounds} rounds in {r['seconds']} s")
    for name in METRICS[:-1]:
        print(f"  {name:10s} round 1 {m[name][0]:.3f}   round {a.rounds} {m[name][-1]:.3f}")
    if a.new_item_round:
        print(f"  new item first shown in round {r['new_item_first_round']}, {sum(m['new_item']):.0f} impressions in all")
    if r["watched"]:
        last = r["watched"]["feeds"][-1]
        print("  the visitor's final feed: " + "; ".join(f"{w.titles[i]} [{TOPICS[int(w.topic[i])]['name']}]" for i in last))
    print(f"wrote {base}.json and {base}.png")


def cmd_compare(a):
    os.makedirs(a.out, exist_ok=True)
    s = settings(a)
    results = {"seed": a.seed, "rounds": a.rounds, "users": a.users, "items": a.items, "k": a.k, "settings": s, "policies": {}}
    w = world_for(a, False)
    runs = []
    for pol in policies.POLICIES:
        t0 = time.time()
        r = run(w, pol, a.rounds, **s)
        r["seconds"] = round(time.time() - t0, 2)
        results["policies"][pol] = r
        runs.append((policies.LABELS[pol], r))
        m = r["metrics"]
        print(f"{policies.LABELS[pol]:13s} {r['seconds']:4.1f}s  diversity {m['diversity'][0]:.2f} -> {m['diversity'][-1]:.2f}  coverage {m['coverage'][-1]:.2f}  gini {m['gini'][-1]:.2f}  utility {m['utility'][-1]:.2f}  clicks {m['clicks'][-1]:.2f}  bubble {m['bubble'][-1]:.2f}  top1% {m['top1'][-1]:.2f}")
    plot_metrics(runs, os.path.join(a.out, "policies.png"))

    sw = a.switch_round or max(1, a.rounds // 2)
    swp = a.switch_policy or "ucb"
    r = run(w, "exploit", a.rounds, switch_round=sw, switch_policy=swp, **s)
    results["switch"] = r
    print(f"exploit -> {swp} at round {sw}: coverage {r['metrics']['coverage'][sw - 2]:.2f} before, {r['metrics']['coverage'][-1]:.2f} at the end; gini {r['metrics']['gini'][sw - 2]:.2f} -> {r['metrics']['gini'][-1]:.2f}")
    plot_metrics([("Pure exploit", results["policies"]["exploit"]), (policies.LABELS[swp], results["policies"][swp]),
                  (f"Exploit, then {policies.LABELS[swp]} from round {sw}", r)], os.path.join(a.out, "switch.png"))

    nr = a.new_item_round or max(1, a.rounds // 3)
    wn = world_for(a, True)
    results["new_item"] = {"round": nr, "policies": {}}
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    fig, ax = plt.subplots(figsize=(7, 4.2))
    for pol in policies.POLICIES:
        r = run(wn, pol, a.rounds, new_item_round=nr, **s)
        cum = np.cumsum(r["metrics"]["new_item"])
        results["new_item"]["policies"][pol] = {"first_round": r["new_item_first_round"], "impressions": r["metrics"]["new_item"]}
        ax.plot(range(1, a.rounds + 1), cum, label=policies.LABELS[pol], lw=1.8)
        print(f"new item (quality 0.9, released round {nr}) under {policies.LABELS[pol]:13s}: first shown round {r['new_item_first_round']}, {cum[-1]:.0f} impressions by round {a.rounds}")
    ax.axvline(nr, color="0.6", ls="--", lw=1)
    ax.set_xlabel("round")
    ax.set_ylabel("impressions of the new item so far")
    ax.set_title(f"A good new item released in round {nr}: how fast each policy finds it")
    ax.grid(alpha=0.3)
    ax.legend(fontsize=9)
    fig.tight_layout()
    fig.savefig(os.path.join(a.out, "new_item.png"), dpi=130)
    plt.close(fig)

    # one user's feed, round by round, under three policies
    user = 0
    strips = []
    for pol in ["exploit", "ucb", "mmr"]:
        r = run(w, pol, a.rounds, watch=user, **s)
        strips.append((policies.LABELS[pol], [[int(w.topic[i]) for i in f] for f in r["watched"]["feeds"]]))
    plot_feed_strip(w, strips, user, os.path.join(a.out, "feed_strip.png"))

    with open(os.path.join(a.out, "compare.json"), "w") as f:
        json.dump(results, f, indent=1)
    print(f"wrote {a.out}/compare.json, policies.png, switch.png, new_item.png, feed_strip.png")


def main(argv=None):
    p = argparse.ArgumentParser(prog="python -m feedloop", description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    sub = p.add_subparsers(dest="cmd", required=True)
    pr = sub.add_parser("run", help="run one policy")
    pr.add_argument("--policy", choices=policies.POLICIES, default="exploit")
    add_common(pr)
    pr.set_defaults(fn=cmd_run)
    pc = sub.add_parser("compare", help="run every policy and the two experiments; write the figures")
    add_common(pc)
    pc.set_defaults(fn=cmd_compare)
    a = p.parse_args(argv)
    a.fn(a)


if __name__ == "__main__":
    main(sys.argv[1:])