# 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, 12, 73, 27), 'random.random', 'random.random', ({}, {}), '()', False, 'import random\n'), ((51, 26, 51, 61), 'random.choice', 'random.choice', ({(51, 40, 51, 60): 'exp.datapoints[-10:]'}, {}), '(exp.datapoints[-10:])', False, 'import random\n'), ((76, 20, 76, 41), 'random.uniform', 'random.uniform', ({(76, 35, 76, 36): '(0)', (76, 38, 76, 40): '(10)'}, {}), '(0, 10)', False, 'import random\n'), ((78, 25, 78, 64), 'random.choice', 'random.choice', ({(78, 39, 78, 63): "['sgd', 'adam', 'adamw']"}, {}), "(['sgd', 'adam', 'adamw'])", False, 'import random\n'), ((80, 17, 80, 57), 'random.choice', 'random.choice', ({(80, 31, 80, 56): "['red', 'green', 'black']"}, {}), "(['red', 'green', 'black'])", False, 'import random\n'), ((85, 27, 85, 47), 'random.uniform', 'random.uniform', ({(85, 42, 85, 43): '1', (85, 45, 85, 46): '5'}, {}), '(1, 5)', False, 'import random\n'), ((87, 27, 87, 47), 'random.uniform', 'random.uniform', ({(87, 42, 87, 43): '1', (87, 45, 87, 46): '5'}, {}), '(1, 5)', False, 'import random\n'), ((98, 11, 98, 26), 'random.random', 'random.random', ({}, {}), '()', False, 'import random\n'), ((99, 34, 99, 54), 'random.uniform', 'random.uniform', ({(99, 49, 99, 50): '1', (99, 52, 99, 53): '5'}, {}), '(1, 5)', False, 'import random\n'), ((101, 34, 101, 80), 'random.choice', 'random.choice', ({(101, 48, 101, 79): '[math.inf, -math.inf, math.nan]'}, {}), '([math.inf, -math.inf, math.nan])', False, 'import random\n'), ((107, 27, 107, 49), 'random.uniform', 'random.uniform', ({(107, 42, 107, 44): '(10)', (107, 46, 107, 48): '(90)'}, {}), '(10, 90)', False, 'import random\n'), ((110, 32, 110, 54), 'random.uniform', 'random.uniform', ({(110, 47, 110, 48): '(1)', (110, 50, 110, 53): '(100)'}, {}), '(1, 100)', False, 'import random\n'), ((129, 21, 129, 55), 'random.choice', 'random.choice', ({(129, 35, 129, 54): 'xp.datapoints[-10:]'}, {}), '(xp.datapoints[-10:])', False, 'import random\n'), ((75, 24, 75, 45), 'random.uniform', 'random.uniform', ({(75, 39, 75, 41): '(-5)', (75, 43, 75, 44): '(0)'}, {}), '(-5, 0)', False, 'import random\n'), ((92, 26, 92, 52), 'random.uniform', 'random.uniform', ({(92, 41, 92, 45): '-1.0', (92, 47, 92, 51): '10.0'}, {}), '(-1.0, 10.0)', False, 'import random\n'), ((106, 32, 106, 53), 'random.uniform', 'random.uniform', ({(106, 47, 106, 49): '(-5)', (106, 51, 106, 52): '(0)'}, {}), '(-5, 0)', False, 'import random\n'), ((109, 25, 109, 50), 'random.uniform', 'random.uniform', ({(109, 40, 109, 44): '(-0.2)', (109, 46, 109, 49): '(0.2)'}, {}), '(-0.2, 0.2)', False, 'import random\n'), ((179, 22, 179, 44), 'random.gauss', 'random.gauss', ({(179, 35, 179, 38): '(0.0)', (179, 40, 179, 43): '(1.0)'}, {}), '(0.0, 1.0)', False, 'import random\n'), ((191, 15, 191, 70), 'random.choice', 'random.choice', ({(191, 29, 191, 69): "['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h']"}, {}), "(['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h'])", False, 'import random\n'), ((192, 19, 192, 43), 'random.uniform', 'random.uniform', ({(192, 34, 192, 37): '0.0', (192, 39, 192, 42): '1.0'}, {}), '(0.0, 1.0)', False, 'import random\n'), ((238, 20, 238, 42), 'random.uniform', 'random.uniform', ({(238, 35, 238, 36): '(0)', (238, 38, 238, 41): '(100)'}, {}), '(0, 100)', False, 'import random\n'), ((253, 26, 253, 46), 'random.randint', 'random.randint', ({(253, 41, 253, 42): '(0)', (253, 44, 253, 45): '(5)'}, {}), '(0, 5)', False, 'import random\n'), ((254, 26, 254, 47), 'random.randint', 'random.randint', ({(254, 41, 254, 42): '(0)', (254, 44, 254, 46): '(10)'}, {}), '(0, 10)', False, 'import random\n'), ((255, 26, 255, 47), 'random.randint', 'random.randint', ({(255, 41, 255, 42): '(0)', (255, 44, 255, 46): '(25)'}, {}), '(0, 25)', False, 'import random\n'), ((270, 52, 270, 72), 'random.randint', 'random.randint', ({(270, 67, 270, 68): '(0)', (270, 70, 270, 71): '(5)'}, {}), '(0, 5)', False, 'import random\n'), ((271, 42, 271, 63), 'random.randint', 'random.randint', ({(271, 57, 271, 58): '(0)', (271, 60, 271, 62): '(10)'}, {}), '(0, 10)', False, 'import random\n'), ((272, 46, 272, 67), 'random.randint', 'random.randint', ({(272, 61, 272, 62): '(0)', (272, 64, 272, 66): '(25)'}, {}), '(0, 25)', False, 'import random\n'), ((48, 24, 48, 45), 'random.uniform', 'random.uniform', ({(48, 39, 48, 41): '-5', (48, 43, 48, 44): '5'}, {}), '(-5, 5)', False, 'import random\n'), ((77, 20, 77, 32), 'uuid.uuid4', 'uuid.uuid4', ({}, {}), '()', False, 'import uuid\n'), ((96, 26, 96, 50), 'random.uniform', 'random.uniform', ({(96, 41, 96, 44): '100', (96, 46, 96, 49): '101'}, {}), '(100, 101)', False, 'import random\n'), ((108, 21, 108, 33), 'uuid.uuid4', 'uuid.uuid4', ({}, {}), '()', False, 'import uuid\n'), ((111, 29, 111, 40), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((124, 49, 124, 64), 'random.random', 'random.random', ({}, {}), '()', False, 'import random\n'), ((235, 29, 235, 49), 'random.uniform', 'random.uniform', ({(235, 44, 235, 45): '(0)', (235, 47, 235, 48): '(1)'}, {}), '(0, 1)', False, 'import random\n'), ((239, 28, 239, 49), 'random.uniform', 'random.uniform', ({(239, 43, 239, 44): '(0)', (239, 46, 239, 48): '(10)'}, {}), '(0, 10)', False, 'import random\n'), ((47, 31, 47, 52), 'random.uniform', 'random.uniform', ({(47, 46, 47, 48): '-1', (47, 50, 47, 51): '1'}, {}), '(-1, 1)', False, 'import random\n'), ((94, 32, 94, 52), 'random.uniform', 'random.uniform', ({(94, 47, 94, 48): '(1)', (94, 50, 94, 51): '(2)'}, {}), '(1, 2)', False, 'import random\n'), ((256, 30, 256, 50), 'random.randint', 'random.randint', ({(256, 45, 256, 46): '(0)', (256, 48, 256, 49): '(5)'}, {}), '(0, 5)', False, 'import random\n'), ((257, 30, 257, 51), 'random.randint', 'random.randint', ({(257, 45, 257, 46): '(0)', (257, 48, 257, 50): '(15)'}, {}), '(0, 15)', False, 'import random\n'), ((258, 30, 258, 51), 'random.randint', 'random.randint', ({(258, 45, 258, 46): '(0)', (258, 48, 258, 50): '(25)'}, {}), '(0, 25)', False, 'import random\n'), ((346, 54, 346, 76), 'random.uniform', 'random.uniform', ({(346, 69, 346, 71): '(15)', (346, 73, 346, 75): '(32)'}, {}), '(15, 32)', False, 'import random\n'), ((119, 37, 119, 49), 'uuid.uuid4', 'uuid.uuid4', ({}, {}), '()', False, 'import uuid\n')]