import collections
import random
import time
import typing
import os

import logging
import dataclasses

logging.basicConfig(level=20)

log = logging.getLogger()

log.info("initialize")

import aim.distributed_hugging_face

class Callback(aim.distributed_hugging_face.AimCallback):
    # I use this just to avoid spinning up an AIM Server, it collects the metrics, computes some statistical data,
    # and then dumps those on the disk.
    def on_train_end(self, args, state, control, **kwargs):
        run_metrics = []
        run_hash = None

        if state.is_world_process_zero:
            run_hash = self._run.hash
            run_metrics = [
                (m.name, m.context.to_dict(), m.values.values_list())
                for m in self._run.metrics()
            ]

        super().on_train_end(args, state, control, **kwargs)

        if not state.is_world_process_zero:
            return

        metrics = []

        for name, context, values in run_metrics:
            try:
                len_values = 0
                _sum = 0
                avg = None
                _min = None
                _max = None

                for x in values:
                    if x is None:
                        continue
                    len_values += 1
                    _sum += x

                    if _min is None or _min > x:
                        _min = x
                    if _max is None or _max < x:
                        _max = x

                if len_values > 0:
                    avg = _sum / len_values

                values = {
                    "avg": avg,
                    "max": _max,
                    "min": _min,
                }
            except ValueError:
                # Don't aggregate properties that are weird
                pass

            metrics.append(
                {
                    "name": name,
                    "values": values,
                    "context": context,
                }
            )

        # Standard
        import json

        data = data = json.dumps(
            {
                "run_hash": run_hash,
                "metrics": metrics,
                "world_size": self.distributed_information.world_size,
            },
            indent=2,
        )

        print(data)
        with open("aim_info.json", "w", encoding="utf-8") as f:
            f.write(data)


cb = Callback(
    main_port=8788,
    repo="my_aim",
    # main_host="localhost",
    system_tracking_interval=1,
    
    # This is so that we can test this on a single machine
    workers_only_on_rank_0=False,
)

log.info("train_init " + str(dataclasses.asdict(cb.distributed_information)))


sleep_seconds = 10
stateType = collections.namedtuple("state", field_names=["is_world_process_zero"])
state = stateType(is_world_process_zero=cb.distributed_information.rank == 0)

cb.on_train_begin(args=None, state=state, control=None)

sleep_seconds += random.random() * 30
log.info(
    f"sleep for {sleep_seconds}" + str(dataclasses.asdict(cb.distributed_information))
)
time.sleep(sleep_seconds)

log.info("train_end" + str(dataclasses.asdict(cb.distributed_information)))

cb.on_train_end(args=None, state=state, control=None)
log.info("stopped" + str(dataclasses.asdict(cb.distributed_information)))