# Copyright (c) Facebook, Inc. and its affiliates.
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

import uuid
import random
import math
import time
import typing as t

from . import experiment as hip

# Demos from the README. If one of those is modified, please modify the readme as well


def demo_change_column_properties() -> hip.Experiment:
    data = [{"param": 1, "loss": 10, "hidden_field": "value1", "c": "red"}, {"param": 2, "loss": 5, "hidden_field": "value2", "c": "black"}]
    exp = hip.Experiment.from_iterable(data)
    exp.parameters_definition["c"].colors = {"red": "rgb(255, 0, 0)", "black": "rgb(0, 0, 0)"}
    exp.parameters_definition["loss"].type = hip.ValueType.NUMERIC_LOG
    exp.display_data(hip.Displays.PARALLEL_PLOT).update({
        'hide': ['hidden_field'],   # This column won't appear in the parallel plot
        'order': ['c']              # Column `c` will be displayed first the in parallel plot
    })
    return exp


def demo_basic_usage() -> hip.Experiment:
    data = [{'dropout': 0.1, 'lr': 0.001, 'loss': 10.0, 'optimizer': 'SGD'},
            {'dropout': 0.15, 'lr': 0.01, 'loss': 3.5, 'optimizer': 'Adam'},
            {'dropout': 0.3, 'lr': 0.1, 'loss': 4.5, 'optimizer': 'Adam'}]
    return hip.Experiment.from_iterable(data)


def demo_line_xy() -> hip.Experiment:
    # DEMO_LINE_XY_BEGIN
    exp = hip.Experiment()
    exp.display_data(hip.Displays.XY).update({
        'axis_x': 'generation',
        'axis_y': 'loss',
    })
    for i in range(200):
        dp = hip.Datapoint(
            uid=str(i),
            values={
                'generation': i,
                'param': 10 ** random.uniform(-1, 1),
                'loss': random.uniform(-5, 5),
            })
        if i > 10:
            from_parent = random.choice(exp.datapoints[-10:])
            dp.from_uid = from_parent.uid  # <-- Connect the parent to the child
            dp.values['loss'] += from_parent.values['loss']  # type: ignore
            dp.values['param'] *= from_parent.values['param']  # type: ignore
        exp.datapoints.append(dp)
    # DEMO_LINE_XY_END
    return exp


def demo_bug_uid() -> hip.Experiment:
    return hip.Experiment.from_iterable([{'a': 1, 'b': 2, 'uid': 50.0}, {'a': 2, 'b': 3, 'uid': 49.33}])


def demo(n: int = 100) -> hip.Experiment:
    xp = hip.Experiment()
    xp.display_data(hip.Displays.XY).update({
        'axis_x': 'time',
        'axis_y': 'exp_metric',
    })

    # Some fake PBT-ish data
    def fake_params() -> t.Dict[str, hip.DisplayableType]:
        r = random.random()
        p: t.Dict[str, hip.DisplayableType] = {
            "lr": 10 ** random.uniform(-5, 0),
            "seed": random.uniform(0, 10),
            "name": uuid.uuid4().hex[:6],
            "optimizer": random.choice(["sgd", "adam", "adamw"]),
            "r": r,
            "c": random.choice(["red", "green", "black"]),
        }
        if r < 0.1:
            del p['optimizer']
        if r > 0.3:
            p["optionA"] = random.uniform(1, 5)
        else:
            p["optionB"] = random.uniform(1, 5)

        if r < 0.2:
            p["pctile"] = -1.0
        elif r < 0.5:
            p["pctile"] = random.uniform(-1.0, 10.0)
        elif r < 0.8:
            p["pctile"] = 10 ** random.uniform(1, 2)
        else:
            p["pctile"] = random.uniform(100, 101)

        if random.random() > 0.3:
            p["special_values"] = random.uniform(1, 5)
        else:
            p["special_values"] = random.choice([math.inf, -math.inf, math.nan])
        return p

    def fake_metrics(tm: float) -> t.Dict[str, hip.DisplayableType]:
        return {
            "exp_metric": 10 ** random.uniform(-5, 0),
            "pct_success": random.uniform(10, 90),
            "chkpt": uuid.uuid4().hex[:6],
            "time": tm + random.uniform(-0.2, 0.2),
            "force_numericlog": random.uniform(1, 100),
            'timestamp': int(time.time() + (task_idx * 2000)),
        }

    current_pop: t.List[t.Dict[str, t.Any]] = [dict(uid=f"init{i}", params=fake_params(), last_ckpt_uid=None) for i in range(10)]
    continue_num = 0
    for task_idx in range(n):
        # All drop checkpoints
        for p in current_pop:
            ckpt_uid = f"{p['uid']}_{uuid.uuid4().hex[:6]}"
            xp.datapoints.append(hip.Datapoint(uid=ckpt_uid, from_uid=p['last_ckpt_uid'], values={**p['params'], **fake_metrics(task_idx)}))
            p['last_ckpt_uid'] = ckpt_uid

        # Randomly drop some
        current_pop = [p for p in current_pop if random.random() > 0.3]

        # Respawn as needed
        for _ in range(10 - len(current_pop)):
            continue_num += 1
            parent = random.choice(xp.datapoints[-10:])
            current_pop.append(dict(uid=f"continue{continue_num}", params=fake_params(), last_ckpt_uid=parent.uid))
    xp.parameters_definition["c"].colors = {"red": "rgb(255, 0, 0)", "green": "rgb(0, 255, 0)", "black": "rgb(0, 0, 0)"}
    xp.parameters_definition["force_numericlog"].type = hip.ValueType.NUMERIC_LOG
    xp.parameters_definition["pctile"].type = hip.ValueType.NUMERIC_PERCENTILE
    xp.parameters_definition["timestamp"].type = hip.ValueType.TIMESTAMP
    return xp


def demo_customize() -> hip.Experiment:
    exp = demo()
    # EXPERIMENT_SETTINGS_SNIPPET2_BEGIN
    # Provide configuration for the parallel plot
    exp.display_data(hip.Displays.PARALLEL_PLOT).update({
        # Hide some columns in the parallel plot
        'hide': ['optionB'],
        # Specify the order for others
        'order': ['time'],  # Put column time first on the left
    })

    # Provide configuration for the table with all the rows
    exp.display_data(hip.Displays.TABLE).update({
        # Don't display `uid` and `from_uid` columns to the user
        'hide': ['uid', 'from_uid'],
        # In the table, order rows by default
        'order_by': [['pct_success', 'desc']],
        # Specify the order for columns
        'order': ['time'],  # Put column time first on the left
    })

    # Provide configuration for the XY graph
    exp.display_data(hip.Displays.XY).update({
        # Default X axis for the XY plot
        'axis_x': 'time',
        # Default Y axis
        'axis_y': 'lr',
        # Configure lines
        'lines_thickness': 1.0,
        'lines_opacity': 0.1,
        # Configure dots
        'dots_thickness': 2.0,
        'dots_opacity': 0.3,
    })
    # EXPERIMENT_SETTINGS_SNIPPET2_END
    return exp


def demo_force_scale() -> hip.Experiment:
    xp = hip.Experiment()
    for _ in range(100):
        values = [abs(random.gauss(0.0, 1.0)) for _ in range(4)]
        xp.datapoints.append(hip.Datapoint({
            f"value{i}": v / sum(values)
            for i, v in enumerate(values)
        }))
    for i in range(4):
        xp.parameters_definition[f"value{i}"].force_range(0.0, 1.0)
    return xp


def demo_distribution(**kwargs: t.Any) -> hip.Experiment:
    xp = hip.Experiment.from_iterable([{
        'cat': random.choice(["a", "b", "c", "d", "e", "f", "g", "h"]),
        'numeric': random.uniform(0.0, 1.0),
    } for i in range(1000)])
    xp.display_data(hip.Displays.DISTRIBUTION).update(kwargs)
    return xp


def demo_bool() -> hip.Experiment:
    return hip.Experiment.from_iterable([
        {"bool": True},
        {"bool": False}
    ])


def demo_color_interpolate() -> hip.Experiment:
    exp = demo()
    exp.parameters_definition["exp_metric"].colormap = "interpolateSinebow"
    return exp


def demo_color_scheme_ylrd() -> hip.Experiment:
    exp = demo()
    exp.parameters_definition["exp_metric"].colormap = "schemeYlOrRd"
    return exp


def demo_color_scheme_accent() -> hip.Experiment:
    exp = demo()
    exp.parameters_definition["exp_metric"].colormap = "schemeAccent"
    return exp


def demo_color_interpolate_inverse() -> hip.Experiment:
    exp = demo_color_interpolate()
    assert exp.parameters_definition["exp_metric"].colormap is not None
    exp.parameters_definition["exp_metric"].colormap += "#inverse"
    return exp


def demo_axis_style() -> hip.Experiment:
    data: t.List[t.Dict[str, t.Any]] = []
    for _ in range(100):
        data.append({
            **{
                f'param{i}': random.uniform(0, 1)
                for i in range(6)
            },
            'loss': random.uniform(0, 100),
            'metric': 10 ** random.uniform(0, 10)
        })
    xp = hip.Experiment.from_iterable(data)
    for i in range(6):
        xp.parameters_definition[f"param{i}"].label_css = "badge badge-pill badge-secondary"
    xp.parameters_definition["loss"].label_css = "badge badge-pill badge-primary"
    xp.parameters_definition["metric"].label_css = "badge badge-pill badge-info"
    return xp


def demo_categorical() -> hip.Experiment:
    data: t.List[t.Dict[str, t.Any]] = []
    for _ in range(100):
        data.append({
            'cat_num_05': random.randint(0, 5),
            'cat_num_15': random.randint(0, 10),
            'cat_num_25': random.randint(0, 25),
            'cat_str_05': f's{random.randint(0, 5)}',
            'cat_str_15': f's{random.randint(0, 15)}',
            'cat_str_25': f's{random.randint(0, 25)}',
        })
    xp = hip.Experiment.from_iterable(data)
    for param in ["cat_num_05", "cat_num_15", "cat_num_25"]:
        xp.parameters_definition[param].type = hip.ValueType.CATEGORICAL
    xp.colorby = 'cat_num_25'
    return xp


def demo_long_names() -> hip.Experiment:
    return hip.Experiment.from_iterable([
        {
            'some very very long name for a field': random.randint(0, 5),
            'this one is also very long': random.randint(0, 10),
            'another.long.one.but.with.dots': random.randint(0, 25),
        }
        for _ in range(100)
    ])


def demo_force_constant_pplot() -> hip.Experiment:
    exp = hip.Experiment.from_iterable([
        {'uid': 123, 'a': 1, 'b': 3},
        {'uid': 345, 'a': 2, 'b': 3}
    ])
    exp.parameters_definition["b"].force_range(0, 100)
    return exp


def demo_first_value_nan() -> hip.Experiment:
    return hip.Experiment.from_iterable([
        {},
        {'a': None},
        {'a': 2},
        {'a': 2.1},
        {'a': 2.2},
        {'a': 5.5},
        {'a': math.nan},
    ])


def demo_weighted_rows() -> hip.Experiment:
    experiment = hip.Experiment.from_iterable([
        {'w': 1.0, 'a': 1, 'b': 1},
        {'w': 2.0, 'a': 2, 'b': 1},
        {'w': -2.0, 'a': 2, 'b': 1},
        {'w': math.inf, 'a': 2, 'b': 2},
        {'w': 'not_a_number', 'a': 2, 'b': 3},
        {'w': None, 'a': 3, 'b': 3},
        {'a': 4, 'b': 3},
    ])
    experiment.weightcolumn = "w"
    return experiment


def demo_3xcols() -> hip.Experiment:
    xp = demo()
    for i in range(2):
        new_xp = demo()
        for dp, new_dp in zip(xp.datapoints, new_xp.datapoints):
            dp.values.update({
                f"{k}{i}": v
                for k, v in new_dp.values.items()
            })
    return xp


def demo_col_html() -> hip.Experiment:
    COL1 = "<h1>col1</h1>"
    COL2 = "col_2"
    experiment = hip.Experiment.from_iterable([
        {COL1: 1.0, COL2: 1},
        {COL1: 2.0, COL2: 2},
        {COL1: 3.0, COL2: 3},
    ])
    experiment.parameters_definition[COL2].label_html = "col<sub>2</sub>"
    return experiment


def demo_disable_table() -> hip.Experiment:
    experiment = demo()
    experiment.enabledDisplays.remove(hip.Displays.TABLE)
    return experiment


def demo_big_floats() -> hip.Experiment:
    return hip.Experiment.from_iterable(
        {
            'bigfloat': math.nan if i < 10 else 10 ** random.uniform(15, 32),
        }
        for i in range(100)
    )


README_DEMOS: t.Dict[str, t.Callable[[], hip.Experiment]] = {
    "demo": demo,
    "demo_3xcols": demo_3xcols,
    "demo_big": lambda: demo(1000),
    "demo_change_column_properties": demo_change_column_properties,
    "demo_basic_usage": demo_basic_usage,
    "demo_line_xy": demo_line_xy,
    "demo_bug_uid": demo_bug_uid,
    "demo_force_scale": demo_force_scale,
    "demo_distribution_cat": lambda: demo_distribution(axis="cat"),
    "demo_distribution_num": lambda: demo_distribution(axis="numeric"),
    "demo_distribution_num_100bins": lambda: demo_distribution(axis="numeric", nbins=100),
    "demo_bool": demo_bool,
    "demo_color_interpolate": demo_color_interpolate,
    "demo_color_scheme_ylrd": demo_color_scheme_ylrd,
    "demo_color_scheme_accent": demo_color_scheme_accent,
    "demo_axis_style": demo_axis_style,
    "demo_categorical": demo_categorical,
    "demo_customize": demo_customize,
    "demo_long_names": demo_long_names,
    "demo_force_constant_pplot": demo_force_constant_pplot,
    "demo_color_interpolate_inverse": demo_color_interpolate_inverse,
    "demo_first_value_nan": demo_first_value_nan,
    "demo_weighted_rows": demo_weighted_rows,
    "demo_col_html": demo_col_html,
    "demo_disable_table": demo_disable_table,
    "demo_big_floats": demo_big_floats,
}

[(73, 'random.random', 'random.random', False, 'import random\n'), (51, 'random.choice', 'random.choice', False, 'import random\n'), (76, 'random.uniform', 'random.uniform', False, 'import random\n'), (78, 'random.choice', 'random.choice', False, 'import random\n'), (80, 'random.choice', 'random.choice', False, 'import random\n'), (85, 'random.uniform', 'random.uniform', False, 'import random\n'), (87, 'random.uniform', 'random.uniform', False, 'import random\n'), (98, 'random.random', 'random.random', False, 'import random\n'), (99, 'random.uniform', 'random.uniform', False, 'import random\n'), (101, 'random.choice', 'random.choice', False, 'import random\n'), (107, 'random.uniform', 'random.uniform', False, 'import random\n'), (110, 'random.uniform', 'random.uniform', False, 'import random\n'), (129, 'random.choice', 'random.choice', False, 'import random\n'), (75, 'random.uniform', 'random.uniform', False, 'import random\n'), (92, 'random.uniform', 'random.uniform', False, 'import random\n'), (106, 'random.uniform', 'random.uniform', False, 'import random\n'), (109, 'random.uniform', 'random.uniform', False, 'import random\n'), (179, 'random.gauss', 'random.gauss', False, 'import random\n'), (191, 'random.choice', 'random.choice', False, 'import random\n'), (192, 'random.uniform', 'random.uniform', False, 'import random\n'), (238, 'random.uniform', 'random.uniform', False, 'import random\n'), (253, 'random.randint', 'random.randint', False, 'import random\n'), (254, 'random.randint', 'random.randint', False, 'import random\n'), (255, 'random.randint', 'random.randint', False, 'import random\n'), (270, 'random.randint', 'random.randint', False, 'import random\n'), (271, 'random.randint', 'random.randint', False, 'import random\n'), (272, 'random.randint', 'random.randint', False, 'import random\n'), (48, 'random.uniform', 'random.uniform', False, 'import random\n'), (77, 'uuid.uuid4', 'uuid.uuid4', False, 'import uuid\n'), (96, 'random.uniform', 'random.uniform', False, 'import random\n'), (108, 'uuid.uuid4', 'uuid.uuid4', False, 'import uuid\n'), (111, 'time.time', 'time.time', False, 'import time\n'), (124, 'random.random', 'random.random', False, 'import random\n'), (235, 'random.uniform', 'random.uniform', False, 'import random\n'), (239, 'random.uniform', 'random.uniform', False, 'import random\n'), (47, 'random.uniform', 'random.uniform', False, 'import random\n'), (94, 'random.uniform', 'random.uniform', False, 'import random\n'), (256, 'random.randint', 'random.randint', False, 'import random\n'), (257, 'random.randint', 'random.randint', False, 'import random\n'), (258, 'random.randint', 'random.randint', False, 'import random\n'), (346, 'random.uniform', 'random.uniform', False, 'import random\n'), (119, 'uuid.uuid4', 'uuid.uuid4', False, 'import uuid\n')]