"""
This model provides
"""

import typing
import jax
import jax.numpy as jnp
import flax
import typing
from flax.training import train_state
from flax import struct
import torch.utils.data as torchdata
from functools import partial
import tqdm
from torch.utils.tensorboard import SummaryWriter

from .progress import tqdm_style

class RunningMean:
    def __init__(self, window_size: int = 100):
        self.window_size = window_size
        self.values = []
        self.mean = None
        self.current_position = 0

    def update(self, value):
        self.values.append(value)
        self.current_position += 1
        if len(self.values) < self.window_size:
            self.mean = None
        elif len(self.values) == self.window_size:
            self.mean = sum(self.values) / self.window_size
        else:
            self.mean = (
                self.mean
                - self.values[self.current_position - self.window_size - 1]
                / self.window_size
                + value / self.window_size
            )
        return

class TrainState(train_state.TrainState):
    """
    Arguments
    ---------

    - key: jax.random.KeyArray:
        The key for the random number generator in the models.

    - batch_stats: typing.Any:
        This is the batch statistics for the batch normalization.

    """

    key: jax.random.KeyArray = None
    batch_stats: typing.Any = None
    train: bool = struct.field(pytree_node=False, default=False)

class TrainStateSourceMapping(train_state.TrainState):
    """
    This is the train state used for the models that allow you
    to caluclate the source mapping each batch.

    Arguments
    ---------

    - predictor_fn: typing.Callable:
        The function that will be used to get the predictions from the model.

    - source_mapper_raw: jax.core.FrozenDict:
        The parameters of the source mapper.

    - weight: float:
        This is the ratio of weight given to the
        new mapper vs the current mapper.
        :code:`current_mapper = (weight/(weight+1))*new_mapper + (1/(weight+1))*current_mapper`.
        Defaults to :code:`1.0`.

    - key: jax.random.KeyArray:
        The key for the random number generator in the models.

    - batch_stats: typing.Any:
        This is the batch statistics for the batch normalization.

    - train: bool:
        Whether to pass :code:`train=False` or :code:`train=True` to the model when
        the outputs are being predicted before doing the source mapping
        calculations. Defaults to :code:`False`.

    """

    predictor_fn: typing.Callable = struct.field(pytree_node=False)
    source_mapper_raw: typing.Any
    weight: float = 1.0
    batch_stats: typing.Any = None
    key: jax.random.KeyArray = None
    train: bool = struct.field(pytree_node=False, default=False)

    def apply_source_mapping(
        self,
        batch,
    ):

        model_params = dict(
            params=self.params,
            source_mapper=self.source_mapper,
            **(
                {"batch_stats": self.batch_stats}
                if self.batch_stats is not None
                else {}
            ),
        )

        inputs, source, label = batch[0], batch[1], batch[2]
        y_hat = self.apply_fn(
            model_params,
            inputs,
            method=self.predictor_fn,
            **({"train": self.train} if self.train is not None else {}),
        )
        new_source_mapper_raw = self.mapper_update(
            y_hat, label, source, weight=self.weight
        )

        return self.replace(source_mapper_raw=new_source_mapper_raw)

    @jax.jit
    def mapper_update(self, y_hat, label, source, weight):
        current_mapper_raw = flax.core.unfreeze(self.source_mapper_raw).copy()

        n_labels = y_hat.shape[-1]
        n_sources = current_mapper_raw["kernel"].shape[0]

        label = jax.nn.one_hot(label, n_labels)
        y_hat_argmax = jnp.argmax(y_hat, axis=-1)

        # finding the mask for which of the data points have the source s and label y
        mask = jnp.logical_and(
            (source == jnp.arange(n_sources).reshape(-1, 1)).T.reshape(
                -1, n_sources, 1
            ),
            (y_hat_argmax == jnp.arange(n_labels).reshape(-1, 1)).T.reshape(
                -1, 1, n_labels
            ),
        )
        # this has shape (BATCH_SIZE, n_sources, n_labels)

        # calculate the mean
        new_mapper_raw = jnp.mean(
            # this has shape (batch_size, n_sources, n_labels, n_labels)
            jnp.stack([jnp.stack([label] * n_labels, axis=1)] * n_sources, axis=1),
            axis=0,
            where=mask.reshape(-1, n_sources, n_labels, 1),
        )

        # where the new mapper is nan, use the current mapper,
        # otherwise scale the current mapper by the weighting.
        # this is the equivalent of
        # mask = ~jnp.isnan(new_mapper)
        # current_mapper_raw[mask] = (
        #     (1/(weight+1))*current_mapper_raw[mask] + (weight/(weight+1))*new_mapper[mask]
        # )
        current_mapper_raw["kernel"] = jnp.where(
            jnp.isnan(new_mapper_raw),
            current_mapper_raw["kernel"],
            (1 / (weight + 1)) * current_mapper_raw["kernel"]
            + (weight / (weight + 1)) * new_mapper_raw,
        )

        return flax.core.freeze(current_mapper_raw)

    @property
    @jax.jit
    def source_mapper(self):
        return flax.core.freeze(
            {
                "kernel": jnp.where(
                    self.source_mapper_raw["kernel"]
                    == jnp.max(
                        self.source_mapper_raw["kernel"], axis=-1, keepdims=True
                    ),
                    jnp.float32(1),
                    jnp.float32(0),
                )
            }
        )

    @classmethod
    def create(cls, *, apply_fn, params, tx, source_mapper, **kwargs):
        """Creates a new instance with `step=0` and initialized `opt_state`."""
        opt_state = tx.init(params)
        return cls(
            step=0,
            apply_fn=apply_fn,
            params=params,
            tx=tx,
            opt_state=opt_state,
            source_mapper_raw=source_mapper,
            **kwargs,
        )

# ==== the loss function ====
def _calc_loss(
    params: flax.core.FrozenDict,
    state: flax.training.train_state.TrainState,
    batch: typing.Tuple[jnp.DeviceArray, ...],
    train: bool,
    criterion: typing.Callable,
    mutable: typing.List[str] = False,
    rngs: typing.Union[None, typing.List[str]] = None,
    apply_fn_kwargs: typing.Dict[str, typing.Any] = {},
):
    """
    Calculates the loss for a batch. This is the default
    function used when training a model using the
    :code:`FitModel` class.
    """
    inputs = batch[:-1]
    targets = batch[-1]

    model_params = dict(params=params)
    if train:
        if rngs is not None:
            new_key = jax.random.fold_in(key=state.key, data=state.step)
            rngs = {
                col: key for col, key in zip(rngs, jax.random.split(new_key, len(rngs)))
            }

    if mutable != False:
        model_params = dict(params=params, **{m: getattr(state, m) for m in mutable})

    outputs = state.apply_fn(
        model_params,
        *inputs,
        train=train,
        mutable=mutable,
        rngs=rngs,
        **apply_fn_kwargs,
    )

    if mutable != False:
        y_hat = outputs[0]
        updates = outputs[1]
    else:
        y_hat = outputs

    loss_value = (
        criterion(y_hat, targets).sum() / targets.shape[0]
    )  # mean loss value over batch

    if mutable != False and train:
        return loss_value, updates

    return loss_value

# ==== the stepping function ====
def _do_step(
    batch: typing.Tuple[jnp.DeviceArray, ...],
    state: flax.training.train_state.TrainState,
    loss_fn: typing.Callable,
    mutable: typing.List[str] = False,
):
    """
    Performs a step using a train state state and batch.
    This is the default function used when training a
    model using the :code:`FitModel` class.
    """

    output, grads = jax.value_and_grad(loss_fn, has_aux=(mutable != False))(
        state.params,
        state=state,
        batch=batch,
    )

    state = state.apply_gradients(grads=grads)

    if mutable != False:
        loss_value = output[0]
        updates = output[1]
        state = state.replace(**updates)
    else:
        loss_value = output

    return state, loss_value

# ==== the epoch function ====
def _do_epoch(
    cls: object,
    train_dl: torchdata.DataLoader,
    state: flax.training.train_state.TrainState,
    step_fn: typing.Callable,
):
    """
    Performs an epoch using a train state and batch.
    This is the default function used when training a
    model using the :code:`FitModel` class.
    """

    cls.tqdm_bar.set_description(f"Epoch {cls.i_epoch+1}/{cls.n_epochs}")
    loss_value_epoch = 0
    # ==== looping on the batches ====
    for ib, batch in enumerate(train_dl):

        state, loss_value = step_fn(
            batch=batch,
            state=state,
        )
        loss_value_epoch += loss_value.item() * len(batch[0])
        cls.running_mean.update(loss_value.item())
        if cls.writer is not None:
            cls.writer.add_scalar(
                "training_loss", loss_value.item(), cls.i_epoch * len(train_dl) + ib
            )
        if cls.running_mean.mean is not None:
            cls.postfix["training loss"] = f"{cls.running_mean.mean:.2f}"
        cls.tqdm_bar.set_postfix(cls.postfix)
        cls.tqdm_bar.update(1)

    loss_value_epoch *= 1 / len(train_dl.dataset)
    cls.loss_dict["train"].append(loss_value_epoch)

    return cls, state

# ==== the eval function ====
def _do_eval(
    cls: object,
    eval_dl: torchdata.DataLoader,
    state: flax.training.train_state.TrainState,
    loss_fn: typing.Callable,
    early_stop: flax.training.early_stopping = None,
    step: int = 0,
):
    """
    Performs a validation pass using a train state
    and batch. This is the default function used when training a
    model using the :code:`FitModel` class.
    """

    # ==== calculating val loss ===
    loss_value_epoch = 0
    for ib, batch in enumerate(eval_dl):
        loss_value = loss_fn(
            params=state.params,
            state=state,
            batch=batch,
        )
        loss_value_epoch += loss_value.item() * len(batch[0])
    loss_value_epoch *= 1 / len(eval_dl.dataset)
    if early_stop is not None:
        _, early_stop = early_stop.update(loss_value_epoch)
    cls.loss_dict["validation"].append(loss_value_epoch)
    cls.postfix["validation loss"] = f"{loss_value_epoch:.2f}"
    cls.tqdm_bar.set_postfix(cls.postfix)
    cls.tqdm_bar.refresh()

    if cls.writer is not None:
        cls.writer.add_scalar("validation_loss", loss_value_epoch, step)
    return cls, early_stop

# ==== the fitting function ====
def _do_fit(
    cls,
    state: flax.training.train_state.TrainState,
    train_dl: torchdata.DataLoader,
    n_epochs: int,
    criterion: typing.Callable,
    val_dl: typing.Union[torchdata.DataLoader, None] = None,
    train_mutable: typing.List[str] = False,
    train_rngs: typing.Union[None, typing.List[str]] = None,
    early_stop: flax.training.early_stopping = None,
    writer: typing.Union[None, SummaryWriter] = None,
    apply_fn_kwargs: typing.Dict[str, typing.Any] = {},
    do_jit: bool = True,
):
    """
    Fit a model using the trainstate and data.
    This is the default function used when training a
    model using the :code:`FitModel` class.
    """

    cls.loss_fn_train = partial(
        cls.loss_fn,
        train=True,
        criterion=criterion,
        mutable=train_mutable,
        rngs=train_rngs,
        apply_fn_kwargs=apply_fn_kwargs,
    )

    cls.loss_fn_train = jax.jit(cls.loss_fn_train) if do_jit else cls.loss_fn_train

    cls.step_fn = partial(
        cls.step_fn,
        loss_fn=cls.loss_fn_train,
        mutable=train_mutable,
    )
    cls.step_fn = jax.jit(cls.step_fn) if do_jit else cls.step_fn

    cls.loss_fn_eval = partial(
        cls.loss_fn,
        train=False,
        criterion=criterion,
        mutable=train_mutable,
        apply_fn_kwargs=apply_fn_kwargs,
    )
    cls.loss_fn_eval = jax.jit(cls.loss_fn_eval) if do_jit else cls.loss_fn_eval

    cls.n_epochs = n_epochs
    cls.tqdm_bar = tqdm.tqdm(
        desc="fit",
        total=len(train_dl) * n_epochs,
        disable=not cls.verbose,
        **tqdm_style,
    )
    cls.postfix = {}
    cls.running_mean = RunningMean(window_size=len(train_dl))
    cls.loss_dict = {"train": [], "validation": []}
    cls.i_epoch = 0
    cls.writer = writer

    while cls.i_epoch < cls.n_epochs:
        cls, state = cls.epoch_fn(
            cls,
            train_dl=train_dl,
            state=state,
            step_fn=cls.step_fn,
        )

        cls, early_stop = cls.eval_fn(
            cls,
            eval_dl=val_dl,
            state=state,
            loss_fn=cls.loss_fn_eval,
            early_stop=early_stop,
            step=(cls.i_epoch + 1) * len(train_dl),
        )

        cls.i_epoch += 1

        if cls.writer is not None:
            cls.writer.flush()

        if early_stop is not None:
            if early_stop.should_stop:
                break

    cls.tqdm_bar.close()

    return state

class FitModel(object):
    def __init__(
        self,
        loss_fn: typing.Callable = None,
        step_fn: typing.Callable = None,
        epoch_fn: typing.Callable = None,
        eval_fn: typing.Callable = None,
        fit_fn: typing.Callable = None,
        verbose: bool = True,
    ):
        """
        A class to fit a model using a train state and data.
        This class allows you to mix and match different loss functions,
        steps, epochs, and evaluation functions. This is useful if you
        want to use a different loss function for training and validation,
        but re-use other parts of the training loop.
        """
        self.loss_fn = _calc_loss if loss_fn is None else loss_fn
        self.step_fn = _do_step if step_fn is None else step_fn
        self.epoch_fn = _do_epoch if epoch_fn is None else epoch_fn
        self.eval_fn = _do_eval if eval_fn is None else eval_fn
        self.fit_fn = _do_fit if fit_fn is None else fit_fn

        self.verbose = verbose

        return

    def fit(self, *args, **kwargs):
        return self.fit_fn(self, *args, **kwargs)

### training functions specific to the calculating source mapper models ###

# ==== the loss function ====
def _calc_loss_source_mapper(
    params: flax.core.FrozenDict,
    state: flax.training.train_state.TrainState,
    batch: typing.Tuple[jnp.DeviceArray, ...],
    train: bool,
    criterion: typing.Callable,
    mutable: typing.List[str] = False,
    rngs: typing.Union[None, typing.List[str]] = None,
    apply_fn_kwargs: typing.Dict[str, typing.Any] = {},
    other_params: typing.Dict[str, typing.Any] = {},
):
    inputs = batch[:-1]
    targets = batch[-1]

    model_params = dict(params=params, **other_params)
    if train:
        if rngs is not None:
            new_key = jax.random.fold_in(key=state.key, data=state.step)
            rngs = {
                col: key for col, key in zip(rngs, jax.random.split(new_key, len(rngs)))
            }

    if mutable != False:
        model_params = dict(
            **model_params,
            **{m: getattr(state, m) for m in mutable},
        )

    outputs = state.apply_fn(
        model_params,
        *inputs,
        train=train,
        mutable=mutable,
        rngs=rngs,
        **apply_fn_kwargs,
    )

    if mutable != False:
        y_hat = outputs[0]
        updates = outputs[1]
    else:
        y_hat = outputs

    loss_value = (
        criterion(y_hat, targets).sum() / targets.shape[0]
    )  # mean loss value over batch

    if mutable != False and train:
        return loss_value, updates

    return loss_value

# ==== the stepping function ====
def _do_step_source_mapper(
    batch: typing.Tuple[jnp.DeviceArray, ...],
    state: flax.training.train_state.TrainState,
    loss_fn: typing.Callable,
    mutable: typing.List[str] = False,
):

    output, grads = jax.value_and_grad(loss_fn, has_aux=(mutable != False))(
        state.params,
        state=state,
        batch=batch,
        other_params={"source_mapper": state.source_mapper},
    )

    state = state.apply_gradients(grads=grads)

    if mutable != False:
        loss_value = output[0]
        updates = output[1]
        state = state.replace(**updates)
    else:
        loss_value = output

    state = state.apply_source_mapping(batch)

    return state, loss_value

def _do_eval_source_mapper(
    cls: object,
    eval_dl: torchdata.DataLoader,
    state: flax.training.train_state.TrainState,
    loss_fn: typing.Callable,
    early_stop: flax.training.early_stopping = None,
    step: int = 0,
):
    # ==== calculating val loss ===
    loss_value_epoch = 0
    for ib, batch in enumerate(eval_dl):
        loss_value = loss_fn(
            params=state.params,
            state=state,
            batch=batch,
            other_params={"source_mapper": state.source_mapper},
        )
        loss_value_epoch += loss_value.item() * len(batch[0])
    loss_value_epoch *= 1 / len(eval_dl.dataset)
    if early_stop is not None:
        _, early_stop = early_stop.update(loss_value_epoch)
    cls.loss_dict["validation"].append(loss_value_epoch)
    cls.postfix["validation loss"] = f"{loss_value_epoch:.2f}"
    cls.tqdm_bar.set_postfix(cls.postfix)
    cls.tqdm_bar.refresh()

    if cls.writer is not None:
        cls.writer.add_scalar("validation_loss", loss_value_epoch, step)
    return early_stop


from traits.api import Instance, on_trait_change, Bool
from traitsui.api import Item, VGroup
from src.processing.analysis import AnalysisTabularAdapter

#============= standard library imports ========================
#============= local library imports  ==========================
from src.processing.plotters.api import Ideogram, InverseIsochron, Spectrum
from src.processing.series.blanks_editor import BlanksEditor
from src.processing.figures.base_figure import BaseFigure, GraphSelector
# from src.processing.result import Result
from src.helpers.traitsui_shortcuts import instance_item
from src.processing.series.backgrounds_editor import BackgroundsEditor
from src.processing.series.detector_intercalibration_editor import DetectorIntercalibrationEditor

class FigureGraphSelector(GraphSelector):
    show_ideo = Bool(True)
    show_inverse_isochron = Bool(False)
    show_spectrum = Bool(False)

    def _get_selections(self):
        s = super(FigureGraphSelector, self)._get_selections()
        return s + [    Item('show_ideo', label='Ideogram'),
                        Item('show_inverse_isochron', label='Inv. Isochron'),
                        Item('show_spectrum', label='Spectrum')
                        ]

class Figure(BaseFigure):
    ideogram = Instance(Ideogram)
    _cached_ideogram = Instance(Ideogram)
    spectrum = Instance(Spectrum)
    inverse_isochron = Instance(InverseIsochron)

    graph_selector = Instance(FigureGraphSelector, ())
    blanks_editor = Instance(BlanksEditor)
    backgrounds_editor = Instance(BackgroundsEditor)
    detector_intercalibration_editor = Instance(DetectorIntercalibrationEditor)
    nanalyses = -1

    def _refresh(self, graph, analyses, padding):

#        self.result = Result()
#        result = self.result

#        pl = 70 if self.show_series else 40
#        padding = [pl, 10, 0, 30]
        specpadding = padding
        ideopadding = padding
        isopadding = padding
#        seriespadding = padding
        gs = self.graph_selector
        if gs.show_ideo:
            '''
                need to do self.ig for .on_trait_change to work
            '''
#            print self._cached_ideogram
            if self._cached_ideogram is not None:
                self.ideogram = ideo = self._cached_ideogram
            else:
                self.ideogram = ideo = Ideogram(db=self.db)

            ideo.figure = self
            ideo.on_trait_change(self._update_selected_analysis, 'selected_analysis')
            ideo.on_trait_change(self.refresh, 'ideogram_of_means')

            gideo = ideo.build(analyses, padding=ideopadding)
            if gideo:
                gideo, plots = gideo
#                graph.plotcontainer.add(gideo.plotcontainer)
#                graph.plotcontainer.add(gideo)
                graph.add(gideo)
#                graph.plots += plots
            self._cached_ideogram = ideo
#                result.add(ig)
        else:
            if self.ideogram:
                self._cached_ideogram = self.ideogram.clone_traits()

            del self.ideogram

        if gs.show_spectrum:
            spec = self.spectrum
            if spec is None:
                spec = Spectrum()
                self.spectrum = spec
            gspec = spec.build(analyses, padding=specpadding)
            if gspec:
#                graph.plotcontainer.add(gspec.plotcontainer)
                graph.add(gspec.plotcontainer)
        else:
            del self.spectrum

        if gs.show_inverse_isochron:
            isochron = self.inverse_isochron
            if isochron is None:
                isochron = InverseIsochron()
                self.inverse_isochron = isochron
            giso = isochron.build(analyses, padding=isopadding)
            if giso:
                graph.add(giso.plotcontainer)
        else:
            del self.inverse_isochron

        super(Figure, self)._refresh(graph, analyses, padding)

    def load_analyses(self, *args, **kw):
        super(Figure, self).load_analyses(*args, **kw)
        self.blanks_editor = BlanksEditor(db=self.db,
                                          figure=self
                                          )
        isos = self.isotope_keys[:]
        isos.reverse()
        self.blanks_editor.add(isos)

        self.backgrounds_editor = BackgroundsEditor(db=self.db,
                                          figure=self
                                          )
        self.backgrounds_editor.add(isos)

        self.detector_intercalibration_editor = DetectorIntercalibrationEditor(db=self.db,
                                                                               figure=self)
        self.detector_intercalibration_editor.add(['IC'])

#        keys = self.signal_keys
#        bl_keys = [i for i in keys if i.endswith('bl')]
#        bg_keys = [i for i in keys if i.endswith('bg')]

#        self.blank_table_adapter.iso_keys = bl_keys
#        self.background_table_adapter.iso_keys = bg_keys

#===============================================================================
# handlers
#===============================================================================
    @on_trait_change('analyses:age_dirty,graph_selector:show_+')
    def _refresh_graph(self, obj, name, old, new):
#        print name
        self.refresh(caller='figure._refresh_graph')

#===============================================================================
# views
#===============================================================================

    def _get_bottom_group(self):
        grp = super(Figure, self)._get_bottom_group()

#        blanksgrp , ta = self._analyses_table_factroy('Blanks')
#        self.blank_table_adapter = ta
#
#        grp.content.append(blanksgrp)
#
#        backsgrp, ta = self._analyses_table_factroy('Backgrounds')
#        self.background_table_adapter = ta
#        grp.content.append(backsgrp)

        height = 200
        editblankgrp = VGroup(
                              Item('blanks_editor',
                                   height=height,
                                   show_label=False, style='custom'),
#                              visible_when='blanks_editor',
                              label='Edit Blanks')

        editbackgrp = VGroup(
                              Item('backgrounds_editor',
                                   height=height,
                                   show_label=False, style='custom'),
#                              visible_when='backgrounds_editor',
                              label='Edit Backgrounds')
        editdetintergrp = VGroup(
                              Item('detector_intercalibration_editor',
                                   height=height,
                                   show_label=False, style='custom'),
#                              visible_when='detector_intercalibration_editor',
                              label='Edit Det. Intercal.')

        editideogrp = VGroup(
                             instance_item('ideogram', height=height),
                             visible_when='ideogram',
                             label='Edit Ideo.')
        editinvisogrp = VGroup(
                             instance_item('inverse_isochron', height=height),
                             visible_when='inverse_isochron',
                             label='Edit Inv. Isochron')
        editspecgrp = VGroup(
                             instance_item('spectrum', height=height),
                             visible_when='spectrum',
                             label='Edit Spectrum.')

        grp.content.append(editdetintergrp)
        grp.content.append(editblankgrp)
        grp.content.append(editbackgrp)
        grp.content.append(editideogrp)
        grp.content.append(editinvisogrp)
        grp.content.append(editspecgrp)

        return grp

#============= EOF =============================================


#Assingment 8_3

import threading

def EvenList(x):
    Elist = []
    for i in x:
        if i > 1:
            if i % 2 == 0:
                Elist.append(i)
    print("Even list is : ",Elist)
    print("Sum of Even list is : ",sum(Elist))

def OddList(a):
    Olist = []
    for i in a:
        if i % 2 != 0:
            Olist.append(i)
    print("Odd list is : ",Olist)
    print("Sum of Odd list : ",sum(Olist))

def main():
    Data = []
    count = int(input("Enter the size of list : "))
    for i in range(count):
        Data.append(int(input()))

    print("Entered list is : ",Data)

    thread1 = threading.Thread(target = EvenList,args = (Data,), name = "Even List")
    thread2 = threading.Thread(target = OddList,args = (Data,), name = "Odd List")

    thread1.start()
    thread2.start()

if __name__=="__main__":
    main()


import numpy as np
import random
import pickle
import os
import os.path as osp
import cv2
import pyflex
from softgym.envs.cloth_env import ClothEnv
from copy import deepcopy
from softgym.utils.misc import vectorized_range, vectorized_meshgrid
from softgym.utils.pyflex_utils import random_pick_and_place, center_object
from gym import spaces
import time
import datetime
import pdb

class ClothSpreadEnv(ClothEnv):
    def __init__(self, reward_type = "sparse", cached_states_path='rope_configuration_init_states.pkl', **kwargs):
        """
        :param cached_states_path:
        :param num_picker: Number of pickers if the aciton_mode is picker
        :param kwargs:
        """
        # kwargs["num_picker"] = 1
        super().__init__(**kwargs)
        self.get_cached_configs_and_states(cached_states_path, self.num_variations)
        self.prev_covered_area = None  # Should not be used until initialized

        self._num_key_points = 4
        self.area_thresh = 0.025
        self.reward_type = reward_type

        # disable the patch_deprecated_methods during registration
        self._gym_disable_underscore_compat = True

        self._goal_save_dir = "save/cloth_spread/goals/"
        if not osp.exists("save/cloth_spread/goals/"):
            os.makedirs("save/cloth_spread/goals/")

        obs = self.reset()
        self.observation_space = spaces.Dict(dict(
            desired_goal=spaces.Box(-np.inf, np.inf, shape=obs['achieved_goal'].shape, dtype='float32'),
            achieved_goal=spaces.Box(-np.inf, np.inf, shape=obs['achieved_goal'].shape, dtype='float32'),
            observation=spaces.Box(-np.inf, np.inf, shape=obs['observation'].shape, dtype='float32'),
        ))

    def generate_env_variation(self, num_variations=1, vary_cloth_size=True):
        """ Generate initial states. Note: This will also change the current states! """
        max_wait_step = 300  # Maximum number of steps waiting for the cloth to stablize
        stable_vel_threshold = 0.01  # Cloth stable when all particles' vel are smaller than this
        generated_configs, generated_states = [], []
        default_config = self.get_default_config()

        for i in range(num_variations):
            config = deepcopy(default_config)
            self.update_camera(config['camera_name'], config['camera_params'][config['camera_name']])
            if vary_cloth_size:
                cloth_dimx, cloth_dimy = self._sample_cloth_size()
                config['ClothSize'] = [cloth_dimx, cloth_dimy]
            else:
                cloth_dimx, cloth_dimy = config['ClothSize']
            self.set_scene(config)
            self.action_tool.reset([0., -1., 0.])
            pos = pyflex.get_positions().reshape(-1, 4)
            pos[:, :3] -= np.mean(pos, axis=0)[:3]
            if self.action_mode in ['sawyer', 'franka']:  # Take care of the table in robot case
                pos[:, 1] = 0.57
            else:
                pos[:, 1] = 0.005
            pos[:, 3] = 1
            pyflex.set_positions(pos.flatten())
            pyflex.set_velocities(np.zeros_like(pos))
            pyflex.step()

            num_particle = cloth_dimx * cloth_dimy
            pickpoint = random.randint(0, num_particle - 1)
            curr_pos = pyflex.get_positions()
            original_inv_mass = curr_pos[pickpoint * 4 + 3]
            curr_pos[pickpoint * 4 + 3] = 0  # Set the mass of the pickup point to infinity so that it generates enough force to the rest of the cloth
            pickpoint_pos = curr_pos[pickpoint * 4: pickpoint * 4 + 3].copy()  # Pos of the pickup point is fixed to this point
            pickpoint_pos[1] += np.random.random(1) * 0.5 + 0.5
            pyflex.set_positions(curr_pos)

            # Pick up the cloth and wait to stablize
            for j in range(0, max_wait_step):
                curr_pos = pyflex.get_positions()
                curr_vel = pyflex.get_velocities()
                curr_pos[pickpoint * 4: pickpoint * 4 + 3] = pickpoint_pos
                curr_vel[pickpoint * 3: pickpoint * 3 + 3] = [0, 0, 0]
                pyflex.set_positions(curr_pos)
                pyflex.set_velocities(curr_vel)
                pyflex.step()
                if np.alltrue(np.abs(curr_vel) < stable_vel_threshold) and j > 5:
                    break

            # Drop the cloth and wait to stablize
            curr_pos = pyflex.get_positions()
            curr_pos[pickpoint * 4 + 3] = original_inv_mass
            pyflex.set_positions(curr_pos)
            for _ in range(max_wait_step):
                pyflex.step()
                curr_vel = pyflex.get_velocities()
                if np.alltrue(curr_vel < stable_vel_threshold):
                    break

            center_object()

            if self.action_mode == 'sphere' or self.action_mode.startswith('picker'):
                curr_pos = pyflex.get_positions()
                self.action_tool.reset(curr_pos[pickpoint * 4:pickpoint * 4 + 3] + [0., 0.2, 0.])
            generated_configs.append(deepcopy(config))
            generated_states.append(deepcopy(self.get_state()))
            self.current_config = config  # Needed in _set_to_flatten function
            generated_configs[-1]['flatten_area'] = self._set_to_flatten()  # Record the maximum flatten area

            print('config {}: camera params {}, flatten area: {}'.format(i, config['camera_params'], generated_configs[-1]['flatten_area']))

        return generated_configs, generated_states

    def reset(self):
        self.goal = self._sample_goal()

        self.current_config = self.cached_configs[0]
        self.set_scene(self.cached_configs[0], self.cached_init_states[0])
        self.particle_num = pyflex.get_n_particles()
        self.prev_reward = 0.
        self.time_step = 0
        obs = self._reset()
        achieved_goal = obs.reshape(-1, 3)[:self._num_key_points].flatten()

        return {
            "observation": obs.copy(),
            "achieved_goal": achieved_goal.copy(),
            "desired_goal": self.goal.copy()
        }

    def step(self, action):
        """ If record_continuous_video is set to True, will record an image for each sub-step"""
        frames = []
        # process action
        action = np.clip(action, self.action_space.low, self.action_space.high)
        action = action.reshape(-1, 4)
        action[:, 1] = - 10.
        action = action.reshape(-1)
        for i in range(self.action_repeat):
            self._step(action)

        obs = self._get_obs()
        achived_goal = obs.copy().reshape(-1, 3)[:self._num_key_points].flatten()
        achived_goal = self._normalize_points(achived_goal)
        obs = self._normalize_points(obs)

        desired_goal = self.goal
        reward = self.compute_reward(achived_goal, desired_goal, None)

        obs = {
            "observation": obs.copy(),
            "achieved_goal": achived_goal.copy(),
            "desired_goal": desired_goal.copy()
        }
        info = {'is_success': reward >= - self.dist_thresh}

        self.time_step += 1

        done = False
        if self.time_step >= self.horizon:
            done = True
        return obs, reward, done, info

    def render_goal(self, target_w, target_h):
        img = pyflex.render()
        width, height = self.camera_params['default_camera']['width'], self.camera_params['default_camera']['height']
        img = img.reshape(height, width, 4)[::-1, :, :3]  # Need to reverse the height dimension
        img = img[int(0.25 * height):int(0.75 * height), int(0.25 * width):int(0.75 * width)]
        img = cv2.resize(img.astype(np.uint8), (target_w, target_h))
        return img

    def render(self, mode='rgb_array'):
        if mode == 'rgb_array':
            img = pyflex.render()
            width, height = self.camera_params['default_camera']['width'], self.camera_params['default_camera']['height']
            img = img.reshape(height, width, 4)[::-1, :, :3]  # Need to reverse the height dimension
            goal_img = self.goal_img.copy()
            # attach goal patch on the rendered image
            goal_img[:10, :, :] = 0
            goal_img[:, :10, :] = 0
            goal_img[-10:, :, :] = 0
            goal_img[:, -10:, :] = 0
            img[30:230, 30:230] = goal_img
            return img
        elif mode == 'human':
            raise NotImplementedError

    def _normalize_points(self, points):
        input_shape = points.shape
        pos = pyflex.get_positions().reshape(-1,4)
        points = points.reshape(-1,3)
        points[:,[0,2]] -= np.mean(pos[:,[0, 2]], axis = 0, keepdims=True)
        return points.reshape(input_shape)

    def _sample_goal(self):
        # randomize the goal.
        cloth_dimx, cloth_dimy = self._sample_cloth_size()
        self.cached_configs[0]['ClothSize'] = [cloth_dimx, cloth_dimy]
        self.cached_configs[0]['flatten_area'] = self._set_to_flatten()
        center_object()

        # reset scene
        config = self.cached_configs[0]
        init_state = self.cached_init_states[0]
        self.set_scene(config, init_state)
        self.key_point_indices = self._get_key_point_idx()

        # read goal state
        particle_pos = np.array(pyflex.get_positions()).reshape([-1, 4])[:, :3]
        keypoint_pos = particle_pos[self.key_point_indices, :3]
        goal = keypoint_pos.flatten()

        # visualize the goal scene
        if hasattr(self, 'action_tool'):
            self.action_tool.reset([10, 10, 10])
        self.goal_img = self.render_goal(200, 200)
        return goal

    def _set_to_flatten(self):
        # self._get_current_covered_area(pyflex.get_positions().reshape(-))
        cloth_dimx, cloth_dimz = self.get_current_config()['ClothSize']
        N = cloth_dimx * cloth_dimz
        px = np.linspace(0, cloth_dimx * self.cloth_particle_radius, cloth_dimx)
        py = np.linspace(0, cloth_dimz * self.cloth_particle_radius, cloth_dimz)
        xx, yy = np.meshgrid(px, py)
        new_pos = np.empty(shape=(N, 4), dtype=np.float)
        new_pos[:, 0] = xx.flatten()
        new_pos[:, 1] = self.cloth_particle_radius
        new_pos[:, 2] = yy.flatten()
        new_pos[:, 3] = 1.
        new_pos[:, :3] -= np.mean(new_pos[:, :3], axis=0)
        pyflex.set_positions(new_pos.flatten())
        return self._get_current_covered_area(new_pos)

    def _reset(self):
        """ Right now only use one initial state"""
        self.prev_covered_area = self._get_current_covered_area(pyflex.get_positions())
        if hasattr(self, 'action_tool'):
            curr_pos = pyflex.get_positions()
            cx, cy = self._get_center_point(curr_pos)
            self.action_tool.reset([cx, 0.2, cy])
        pyflex.step()
        self.init_covered_area = None
        info = self._get_info()
        self.init_covered_area = info['performance']
        return self._get_obs()

    def _step(self, action):
        self.action_tool.step(action)
        if self.action_mode in ['sawyer', 'franka']:
            pyflex.step(self.action_tool.next_action)
        else:
            pyflex.step()
        return

    def _get_current_covered_area(self, pos):
        """
        Calculate the covered area by taking max x,y cood and min x,y coord, create a discritized grid between the points
        :param pos: Current positions of the particle states
        """
        pos = np.reshape(pos, [-1, 4])
        min_x = np.min(pos[:, 0])
        min_y = np.min(pos[:, 2])
        max_x = np.max(pos[:, 0])
        max_y = np.max(pos[:, 2])
        init = np.array([min_x, min_y])
        span = np.array([max_x - min_x, max_y - min_y]) / 100.
        pos2d = pos[:, [0, 2]]

        offset = pos2d - init
        slotted_x_low = np.maximum(np.round((offset[:, 0] - self.cloth_particle_radius) / span[0]).astype(int), 0)
        slotted_x_high = np.minimum(np.round((offset[:, 0] + self.cloth_particle_radius) / span[0]).astype(int), 100)
        slotted_y_low = np.maximum(np.round((offset[:, 1] - self.cloth_particle_radius) / span[1]).astype(int), 0)
        slotted_y_high = np.minimum(np.round((offset[:, 1] + self.cloth_particle_radius) / span[1]).astype(int), 100)
        # Method 1
        grid = np.zeros(10000)  # Discretization
        listx = vectorized_range(slotted_x_low, slotted_x_high)
        listy = vectorized_range(slotted_y_low, slotted_y_high)
        listxx, listyy = vectorized_meshgrid(listx, listy)
        idx = listxx * 100 + listyy
        idx = np.clip(idx.flatten(), 0, 9999)
        grid[idx] = 1

        return np.sum(grid) * span[0] * span[1]

        # Method 2
        # grid_copy = np.zeros([100, 100])
        # for x_low, x_high, y_low, y_high in zip(slotted_x_low, slotted_x_high, slotted_y_low, slotted_y_high):
        #     grid_copy[x_low:x_high, y_low:y_high] = 1
        # assert np.allclose(grid_copy, grid)
        # return np.sum(grid_copy) * span[0] * span[1]

    def _get_center_point(self, pos):
        pos = np.reshape(pos, [-1, 4])
        min_x = np.min(pos[:, 0])
        min_y = np.min(pos[:, 2])
        max_x = np.max(pos[:, 0])
        max_y = np.max(pos[:, 2])
        return 0.5 * (min_x + max_x), 0.5 * (min_y + max_y)

    def compute_reward(self, achieved_goal, desired_goal, info):
        particle_pos = pyflex.get_positions()
        curr_covered_area = self._get_current_covered_area(particle_pos)
        flatten_area = self.cached_configs[0]['flatten_area'] #flatten_area equals to max covered area

        if self.reward_type == "sparse":
            r = -(((flatten_area - curr_covered_area) / flatten_area) > 0.025).astype(np.float32)
        else:
            r = curr_covered_area
        return r

    # @property
    # def performance_bound(self):
    #     dimx, dimy = self.current_config['ClothSize']
    #     max_area = dimx * self.cloth_particle_radius * dimy * self.cloth_particle_radius
    #     min_p = 0
    #     max_p = max_area
    #     return min_p, max_p

    def _get_info(self):
        # Duplicate of the compute reward function!
        particle_pos = pyflex.get_positions()
        curr_covered_area = self._get_current_covered_area(particle_pos)
        init_covered_area = curr_covered_area if self.init_covered_area is None else self.init_covered_area
        max_covered_area = self.get_current_config()['flatten_area']
        info = {
            'performance': curr_covered_area,
            'normalized_performance': (curr_covered_area - init_covered_area) / (max_covered_area - init_covered_area),
        }
        if 'qpg' in self.action_mode:
            info['total_steps'] = self.action_tool.total_steps
        return info

    def get_picked_particle(self):
        pps = np.ones(shape=self.action_tool.num_picker)  * -1 # -1 means no particles picked
        for i, pp in enumerate(self.action_tool.picked_particles):
            if pp is not None:
                pps[i] = pp
        return pps


import pygame
from pygame.locals import K_BACKSPACE, K_DELETE, K_RETURN, KEYDOWN, NUMEVENTS, SRCALPHA, USEREVENT

from cluequiz.style import *

TEXTINPUTREADY = USEREVENT + (42 % (NUMEVENTS-USEREVENT))

class TextPrompt:
    def __init__(self, font, width, label, value='', max_width=None, placeholder='', userdata=None):
        self.visible = False
        self.x, self.y, self.w, self.h, self.left, self.top = (None, None, None, None, None, None)
        self.bg = ( 31,  31,  31)
        self.fg = (255, 255, 255)
        self.fog = None

        self.font = font
        self.width = width
        self.label = label
        self.rndrd_label = self.font.render(self.label, True, self.fg, self.bg)
        self.value = value
        self.max_width = max_width
        self.placeholder = placeholder
        self.userdata = userdata

        self.render_value()

    def render_value(self):
        self.rndrd_value = self.font.render(self.value if len(self.value) > 0 else self.placeholder, True, self.fg, CLUE_COLOR)

    def set_style(self, bg=( 31,  31,  31), fg=(255, 255, 255)):
        self.bg = bg
        self.fg = fg
        self.fog = None
        self.rndrd_label = self.font.render(self.label, True, self.fg, self.bg)

    def set_userdata(self, userdata):
        self.userdata = userdata

    def show(self):
        self.visible = True

    def reset(self):
        self.visible = False
        self.value = ''
        self.render_value()

    def is_visible(self):
        return self.visible

    def handle(self, event):
        if not self.visible:
            return

        if event.type == KEYDOWN:
            if event.key == K_RETURN:
                pygame.event.post(pygame.event.Event(TEXTINPUTREADY, { 'userdata': self.userdata, 'value': self.value }))
                self.reset()
            elif event.key == K_DELETE:
                self.reset()
            elif event.key == K_BACKSPACE:
                self.value = self.value[:-1] if len(self.value) > 0 else ''
                self.render_value()
            elif not self.max_width or self.font.size(self.value + event.unicode)[0] <= self.max_width:
                self.value = self.value + event.unicode
                self.render_value()

    def update(self, display):
        if not self.visible:
            return

        dw, dh = display.get_size()
        linesize = self.font.get_linesize()
        if self.w == None:
            min_width = max(self.width+2*CELL_PADDING, self.rndrd_label.get_width()) + 2 * CELL_PADDING
            self.w = max(dw//2, min_width)
            self.x = (dw-self.w) // 2
            self.left = self.x + (self.w-min_width)//2
        if self.h == None:
            min_height = 3 * linesize + 4 * CELL_PADDING
            self.h = max(dh//2, min_height)
            self.y = (dh-self.h) // 2
            self.top = self.y + (self.h-min_height)//2
        if self.fog == None:
            self.fog = pygame.Surface(display.get_size(), flags=SRCALPHA)
            self.fog.fill(FOG_COLOR)

        display.blit(self.fog, (0, 0))
        display.fill(self.bg, rect=pygame.Rect(self.x, self.y, self.w, self.h))

        display.blit(self.rndrd_label, self.rndrd_label.get_rect(centerx=dw//2).move(0, self.top))
        display.fill(CLUE_COLOR, rect=pygame.Rect(self.left+CELL_PADDING, self.top+2*linesize, self.width+2*CELL_PADDING, linesize+2*CELL_PADDING))

        display.set_clip(pygame.Rect(self.left+2*CELL_PADDING, self.top+2*linesize+CELL_PADDING, self.width, linesize))
        display.blit(self.rndrd_value, self.rndrd_value.get_rect(centerx=dw//2).move(0, self.top+2*linesize+CELL_PADDING))
        display.set_clip(None)


import sys, os
from collections import OrderedDict, namedtuple

from dotenv import dotenv_values

def _verify_config(config: OrderedDict, variables: list):
    try:
        assert(all(k in variables for k in config.keys()))
    except Exception as e:
        print('config contains unkown values: %s', config.keys() - variables)
        pass
    try:
        assert(all(k in config.keys() for k in variables))
    except:
        raise Exception('invalid config')
    return config

def parse_chain(env_file: str):
    variables = ['ID', 'GAS_PRICE', 'EXPLORER', 'RPC', 'FACTORY', 'ROUTER', 'TOKEN_A', 'TOKEN_B']
    chainConfig = namedtuple('Chain', variables)
    config = _verify_config(dotenv_values(env_file), variables)
    chain = chainConfig(**config)
    return chain

def parse_env():
    variables = ['PRIVATE_KEY', 'NET', 'DEBUG']
    appConfig = namedtuple('App', variables)
    config = _verify_config(dotenv_values('.env'), variables)
    app = appConfig(**config)
    chain = parse_chain(f'.env.{app.NET}')
    root = namedtuple('Root', ['app', 'chain'])
    return root(app, chain)

if __name__ == '__main__':

    cfg = parse_env()


import matplotlib
matplotlib.use('agg')

import matplotlib.pyplot as plt
import numpy as np
import gizmo_analysis as ga
import utilities as ga_ut

import sys

FIRE_elements      = ['h','he','c','n','o','ne','mg','si','s','ca','fe']
FIRE_metals = ['c','n','o','ne','mg','si','s','ca','fe']

#
# wrapper to load data set and set FIRE
# abundance tables
#
def load_with_FIRE(sim_index, wdir, return_model = False, model_Z = 1.0):
    """
    Convenience wrapper to load a dataset and set the appropriate initial
    abundances and FIRE yield model tables to do age tracer postprocessing
    with default FIRE2 yields
    """

    initial_part = ga.io.Read.read_snapshots(['gas'],'index',0,simulation_directory=wdir)
    initial_abundances = {}

    # -- note: np.average isn't necessary since they all should have the same value...
    for e in FIRE_elements:
        initial_abundances[e]    = np.average(initial_part['gas'].prop('massfraction.' + e))
    initial_abundances['metals'] = np.average(initial_part['gas'].prop('massfraction.metals'))

    part = ga.io.Read.read_snapshots(['gas','star'],
                              'index',
                              sim_index,
                              simulation_directory = wdir)

    FIRE_yield_model = ga.agetracers.FIRE2_yields(model_Z = model_Z # yield table metallicity in solar units,
                                            )
    FIRE_yield_table = ga.agetracers.construct_yield_table(FIRE_yield_model, # object with a `yields` function
                                                              part.ageprop.age_bins/1000.0) # in Gyr

    part.set_yield_table(FIRE_yield_table,          # the 2D table
                         FIRE_yield_model.elements  # list of elements we want to be able to post process
                                                # these can be anything included in the yield model
                        )

    # finally, set the initial abundances:
    #      As generated above, this is a dictionary corresponding to the initial
    #      mass fractions of each element (and all metals). If an element is missing,
    #      it is assumed to have an initial mass fraction of 0
    part.set_initial_abundances(initial_abundances)

    if return_model:
        return part, FIRE_yield_model, FIRE_yield_table
    else:
        return part

def compute_error(part, element, particle_type = 'star', filter=True):
    """
    For a single element, compute and bin the error
    """

    if filter:
        select = part[particle_type].prop('metallicity.o') > -3.8
    else:
        select = part[particle_type].prop('metallicity.o') == part[particle_type].prop('metallicity.o')

    error = np.abs(part[particle_type].prop('metallicity.agetracer.'+element)[select] - #+np.log10(0.68) -
                   part[particle_type].prop('metallicity.' + element)[select])

    return error

def bin_error(part, element, particle_type='star',
              amin = 0.00, amax = 3.0, da = 0.001, logbins=False):
    """
    Return a distribution of the error
    """

    error = compute_error(part,element,particle_type)
    print(element, np.min(error), np.max(error))

    if logbins:
        bins = np.arange(np.log10(amin), np.log10(amax), da)
        _error = np.log10(error)
    else:
        bins = np.arange(amin, amax+0.5*da, da)
        _error = error

    hist, bins = np.histogram(_error, bins=bins)

    stats = {'bins' : bins,
             'hist' : hist,
             'cbins': 0.5*(bins[1:]+bins[:-1]),
             'median' : np.percentile(error,50.0),
             'onesigma' : np.percentile(error,68.0),
             'twosigma' : np.percentile(error,95.0),
             'threesigma' : np.percentile(error,99.7)}

    return  stats

def compute_all_errors(part,particle_type='star',logbins=False,amin=0.0,amax=3,da=0.001):
    """

    """

    all_hist  = {}
    all_stats = {}

    for e in FIRE_metals:
        all_stats[e] = bin_error(part,e,particle_type,amin=amin,amax=amax,da=da,logbins=logbins)

    return all_stats

def generate_analysis(runs):

    all_part = {}
    all_data = {}
    for runname in runs.keys():
        all_part[runname] = load_with_FIRE(600,"./" + runname, model_Z = 0.1)
        all_data[runname] = {}
        all_data[runname] = compute_all_errors(all_part[runname],'star')

    fs = 5
    fig,ax = plt.subplots(3,3, sharex=True,sharey=True)
    fig.set_size_inches(3*fs,3*fs)
    fig.subplots_adjust(hspace=0,wspace=0)

    axi, axj = 0,0
    for e in FIRE_metals:
        axindex = (axi,axj)

        for runname in runs.keys():

            ploty = np.cumsum(all_data[runname][e]['hist']*1.0)
            ploty = ploty/ploty[-1]

            ax[axindex].plot(all_data[runname][e]['cbins'], ploty,
                    lw = 3, label = runs[runname])
        ax[axindex].set_ylim(0.0,1.0)
        ax[axindex].set_xlim(0.001, 1.0)
        ax[axindex].semilogx()

        axj = axj + 1
        if axj >= 3:
            axj = 0
            axi = axi + 1

    for i in np.arange(3):
        ax[(i,0)].set_ylabel("Cumulative Histogram")
        ax[(2,0)].set_xlabel("Error [dex]")

    fig.savefig("rate_error_panel.png")
    #
    # now compute the metallicities for each and the difference
    #

    def average_stats(runname, elements):
        one = two = three = 0.0
        for e in elements:
            one = one + all_data[runname][e]['onesigma']
            two = two + all_data[runname][e]['twosigma']
            three = three  + all_data[runname][e]['threesigma']

        n = 1.0 * len(elements)

        return one / n, two / n, three / n

    f = open('rates_results.dat','w')

    f.write("name alpha_one alpha_two alpha_three wind_one wind_two wind_three ia_one ia_two ia_three\n")

    for run in runs.keys():

        f.write(run)

        for elist in [ ['o','mg','si','ca'],  ['c','n'], ['fe']]:
            one, two, three = average_stats(run, elist)
            f.write(" %.4f %.4f %.4f"%(one,two,three))
        f.write("\n")

    f.close()

def compare_distributions(runs):

    all_part = {}
    all_data = {}

    amin = 0.01
    amax = 2.0
    da = 0.01

    for runname in runs.keys():
        all_part[runname] = load_with_FIRE(600,"./" + runname, model_Z = 0.1)
        all_data[runname] = {}
        all_data[runname] = compute_all_errors(all_part[runname],'star',
                                               amin = amin, amax = amax, da = da,
                                               logbins=True)

    fs = 6
    fig,ax = plt.subplots(1,3, sharex=True,sharey=True)
    fig.set_size_inches(3*fs,1*fs)
    fig.subplots_adjust(wspace=0)

    for axi, element_list in enumerate([ ['o','mg','si','ca'], ['c','n'], ['fe']]):

        for runname in runs.keys():
            # compute average
            avg_hist = np.zeros(np.size(all_data[runname]['c']['hist']))
            count = 0

            for e in element_list:
                avg_hist += all_data[runname][e]['hist']
                count = count + 1
            avg_hist = avg_hist / (1.0*count)

            dz = all_data[runname][e]['bins'][1:] - all_data[runname][e]['bins'][:-1]

            ax[axi].plot(all_data[runname][e]['cbins'],
                         avg_hist / (1.0*np.sum(avg_hist)) / dz,
                         lw = 3, label = runs[runname])

        ax[axi].set_ylim(0,7)
        ax[axi].set_xlim(np.log10(amin),np.log10(amax))
#        ax[axi].semilogx()
        ax[axi].set_xlabel(r"log$_{10}$(Error [dex])")

    ax[0].set_ylabel("dN/d(log(Error))")

    xy = (0.8,0.1)
    ax[0].annotate(r'$\alpha$', xy, xy, xycoords='axes fraction')
    ax[1].annotate("C+N", xy, xy, xycoords='axes fraction')
    ax[2].annotate("Fe", xy, xy, xycoords='axes fraction')

    ax[0].legend(loc='best')

    fig.savefig('distributions.png', bbox_inches='tight', pad_inches=0.0)

    return

def compare_runtime(runs):

    return

def plot_sigma(runs):

    data = np.genfromtxt("rates_results.dat",names=True)

    xvals = np.arange(6)

    fig, ax = plt.subplots(1,3, sharey=True)
    fig.set_size_inches(18,7)
    fig.subplots_adjust(wspace=0)

    colors = {'one': 'C0', 'two' : 'C1', 'three' :'C2'}
    labels = {'one': r'1 $\sigma$', 'two': r'2 $\sigma$', 'three' : r'3 $\sigma$'}

    annotation = {'alpha': r"$\alpha$ (CCSNe)", 'wind' : "C + N (Winds)", 'ia' : "Fe (Ia's)"}

    for i, metal in enumerate(['alpha','wind','ia']):
        for sigma in ['one','two','three']:
            ax[i].plot(xvals, data[metal + '_' + sigma], color = colors[sigma], lw = '2', ls = ':')
            ax[i].scatter(xvals, data[metal + '_' + sigma], s = 30, color = colors[sigma], label = labels[sigma])

        ax[i].set_ylim(0.0,1.0)
        ax[i].set_xlim(-0.5, 5.5)
        ax[i].set_xticks(xvals)
        ax[i].set_xticklabels( list(runs.values()), rotation = 'vertical', fontsize = 12)

        xy = (0.8,0.05)
        ax[i].annotate(annotation[metal], xy,xy,xycoords='axes fraction')

    ax[0].set_ylabel("Simulation - AgeTracer [dex]")
    ax[0].legend(loc='best')

    fig.savefig("sigma_comparison.png", bbox_inches='tight', pad_inches=0.05)

    return

if __name__ == "__main__":

    runs = {"test0" : "R = 1.0", 'test1' : "R = 0.5",
            "test2" : "R = 0.1", "test3" : "R = -1",
            'test4' : "R = -10", "test5" : "R = -100"}

    generate = False
    compare_runtime = False
    compare_sigma = False
    plot_distributions = False
    if len(sys.argv) > 1:
        if 'generate' in sys.argv:
            generate = True

        if 'runtime' in sys.argv:
            compare_runtime = True

        if 'compare_sigma' in sys.argv:
            compare_sigma = True

        if 'plot_distributions' in sys.argv:
            plot_distributions = True

    else:
        generate = True

    if generate:
        generate_analysis(runs)

    if compare_runtime:
        compare_runtime(runs)

    if compare_sigma:
        plot_sigma(runs)

    if plot_distributions:
        compare_distributions(runs)


import datetime
import os
from src.video import Video

import isodate

from googleapiclient.discovery import build

class PlayList:
    __api_key: str = os.getenv('YT_API_KEY')
    __youtube = build('youtube', 'v3', developerKey=__api_key)

    def __init__(self, playlist_id):
        self.__playlist_id = playlist_id

        youtube = self.get_service()
        playlists = youtube.playlists().list(id=self.__playlist_id,
                                             part='contentDetails,snippet').execute()
        self.title = playlists['items'][0]['snippet']['title']
        self.url = 'https://www.youtube.com/playlist?list=' + playlists['items'][0]['id']
        self.video_info = []  # список для информации о видео по его id
        self.get_video_info()

    def get_video_info(self):
        """Функция для получения информации о видео по его id"""

        youtube = self.get_service()
        # получаем информацию плейлиста по его id
        playlist_videos = youtube.playlistItems().list(playlistId=self.__playlist_id,
                                                       part='contentDetails',
                                                       maxResults=50,
                                                       ).execute()

        # формируем список из id всех видеороликов из плейлиста
        video_ids: list[str] = [video['contentDetails']['videoId'] for video in playlist_videos['items']]

        # формируем словарь с информацией о видеороликах из списка "video_ids"
        video_response = youtube.videos().list(part='contentDetails,statistics',
                                               id=','.join(video_ids)).execute()

        # добавляем в список "video_info" информацию о видеороликах из "video_response"
        for video in video_response['items']:
            video_obj = Video(video['id'])
            self.video_info.append({
                'id': video['id'],
                'duration': isodate.parse_duration(video['contentDetails']['duration']),
                'like_count': int(video_obj.like_count),
                'url': video_obj.video_url,
            })

    @classmethod
    def get_service(cls):
        """Возвращает объект для работы с YouTube API"""
        return cls.__youtube

    @property
    def total_duration(self):
        total = datetime.timedelta(seconds=0)
        for video in self.video_info:
            total += video['duration']
        return total

    def show_best_video(self):
        best_video_url = ''
        max_like = 0
        for video in self.video_info:
            if video['like_count'] > max_like:
                max_like = video['like_count']
                best_video_url = video['url']
        return best_video_url


# Licensed under a 3-clause BSD style license - see LICENSE.rst

import numpy as np

from sofia_redux.instruments.hawc.datafits import DataFits
from sofia_redux.instruments.hawc.steps import basehawc
from sofia_redux.instruments.hawc.steps.stepprepare import StepPrepare
from sofia_redux.instruments.hawc.steps.stepdemodulate import StepDemodulate
from sofia_redux.instruments.hawc.tests.resources \
    import DRPTestCase, basic_raw_data

class TestBaseHawc(DRPTestCase):
    def test_readfunc(self, tmpdir):
        """
        Test all the read functions in basehawc.

        All tested together to avoid redoing set-up steps.
        """
        # test data -- increase frames for this test case
        hdul = basic_raw_data(nframe=320, smplfreq=1)
        ffile = str(tmpdir.join('test.fits'))
        hdul.writeto(ffile, overwrite=True)
        inp = DataFits(ffile)
        prep = StepPrepare()(inp)

        # prepare to call basehawc
        step = StepDemodulate()
        step.datain = prep
        step.runstart(prep, {})
        step.readdata()

        # test readchop

        # chop state is potentially made of 0, 1, -1 -- this
        # data set should be evenly split between 1 and -1
        chopstate = basehawc.readchop(step, step.praw['nsamp'], 5.0,
                                      chopflag=True)
        zero = chopstate[chopstate == 0]
        high = chopstate[chopstate == 1]
        low = chopstate[chopstate == -1]
        assert zero.size + high.size + low.size == chopstate.size
        assert high.size == low.size
        assert zero.size == 0

        # when chopflag is False, should be only ones
        chopstate = basehawc.readchop(step, step.praw['nsamp'], 5.0,
                                      chopflag=False)
        assert np.all(chopstate == 1)

        # test readnod

        # nod state is just like chop state, but
        # zeros between nod values for these data
        nodstate = basehawc.readnod(step, step.praw['nsamp'], 5.0,
                                    nodflag=True)
        zero = nodstate[nodstate == 0]
        high = nodstate[nodstate == 1]
        low = nodstate[nodstate == -1]
        assert zero.size + high.size + low.size == nodstate.size
        assert high.size == low.size
        assert zero.size == 16

        # when nodflag is False, should be only ones
        nodstate = basehawc.readnod(step, step.praw['nsamp'], 5.0,
                                    nodflag=False)
        assert np.all(nodstate == 1)

        # test readhwp
        # angles in test data are: 5.0 50.0 27.0 72.0, ~80 samples each

        # hwp state for these data should be mostly 1s, a few 0s
        # (3 at the beginning + 3, between angles);
        # nhwp should be 4
        hwpstate, nhwp = basehawc.readhwp(step, step.praw['nsamp'], 5.0, 10.2)
        assert nhwp == 4
        zero = hwpstate[hwpstate == 0]
        high = hwpstate[hwpstate == 1]
        assert zero.size + high.size == hwpstate.size
        assert zero.size == 6

        # set the tolerance too low -- all invalid points
        hwpstate, nhwp = basehawc.readhwp(step, step.praw['nsamp'], 0, 10.2)
        assert np.all(hwpstate == 0)

        # set a bad data region in the middle of the first angle --
        # will trim data to start after bad region
        step.praw['HWP Angle'][7:20] = 50.0
        hwpstate, nhwp = basehawc.readhwp(step, step.praw['nsamp'], 5.0, 10.2)
        assert np.where(hwpstate == 1)[0][0] == 21
        # hwp index set to -1 for invalid, 0-3 for hwp angles 1-4
        assert step.praw['HWP Index'][20] == -1
        assert step.praw['HWP Index'][21] == 0
        assert step.praw['HWP Index'][-1] == 3

        # set a few bad samples (> tol) at the beginnings of hwp regions
        # should be trimmed to exclude

        # end of first angle
        step.praw['HWP Angle'][78] = 6.5
        step.praw['HWP Angle'][79] = 7.5
        # beginning of third angle
        step.praw['HWP Angle'][160:180] = 24.5
        step.praw['HWP Angle'][180] = 25.5
        hwpstate, nhwp = basehawc.readhwp(step, step.praw['nsamp'], 2.0, 10.2)
        assert np.all(hwpstate[78:80]) == 0
        assert np.all(hwpstate[160:181] == 0)
        assert np.where(step.praw['HWP Index'] == 2)[0][0] == 181


"""icons.py - Load MComix specific icons."""

import gtk
from pkg_resources import resource_string

from mcomix import image_tools
from mcomix import log

def mcomix_icons():
    """ Returns a list of differently sized pixbufs for the
    application icon. """

    sizes = ('16x16', '32x32', '48x48')
    pixbufs = [
        image_tools.load_pixbuf_data(
            resource_string('mcomix.images', size + '/mcomix.png')
        ) for size in sizes
    ]

    return pixbufs

def load_icons():
    _icons = (('gimp-flip-horizontal.png',   'mcomix-flip-horizontal'),
              ('gimp-flip-vertical.png',     'mcomix-flip-vertical'),
              ('gimp-rotate-180.png',        'mcomix-rotate-180'),
              ('gimp-rotate-270.png',        'mcomix-rotate-270'),
              ('gimp-rotate-90.png',         'mcomix-rotate-90'),
              ('gimp-thumbnails.png',        'mcomix-thumbnails'),
              ('gimp-transform.png',         'mcomix-transform'),
              ('tango-enhance-image.png',    'mcomix-enhance-image'),
              ('tango-add-bookmark.png',     'mcomix-add-bookmark'),
              ('tango-archive.png',          'mcomix-archive'),
              ('tango-image.png',            'mcomix-image'),
              ('library.png',                'mcomix-library'),
              ('comments.png',               'mcomix-comments'),
              ('zoom.png',                   'mcomix-zoom'),
              ('lens.png',                   'mcomix-lens'),
              ('double-page.png',            'mcomix-double-page'),
              ('manga.png',                  'mcomix-manga'),
              ('fitbest.png',                'mcomix-fitbest'),
              ('fitwidth.png',               'mcomix-fitwidth'),
              ('fitheight.png',              'mcomix-fitheight'),
              ('fitmanual.png',              'mcomix-fitmanual'),
              ('goto-first-page.png',        'mcomix-goto-first-page'),
              ('goto-last-page.png',         'mcomix-goto-last-page'),
              ('next-page.png',              'mcomix-next-page'),
              ('previous-page.png',          'mcomix-previous-page'),
              ('next-archive.png',           'mcomix-next-archive'),
              ('previous-archive.png',       'mcomix-previous-archive'),
              ('next-directory.png',         'mcomix-next-directory'),
              ('previous-directory.png',     'mcomix-previous-directory'))

    # Load window title icons.
    pixbufs = mcomix_icons()
    gtk.window_set_default_icon_list(*pixbufs)
    # Load application icons.
    factory = gtk.IconFactory()
    for filename, stockid in _icons:
        try:
            icon_data = resource_string('mcomix.images', filename)
            pixbuf = image_tools.load_pixbuf_data(icon_data)
            iconset = gtk.IconSet(pixbuf)
            factory.add(stockid, iconset)
        except Exception:
            log.warning(_('! Could not load icon "%s"'), filename)
    factory.add_default()

# vim: expandtab:sw=4:ts=4


from os.path import dirname, realpath
from typing import List
from itertools import chain
from re import match, findall

class TicketValidator:
    rules = {}

    def addRule(self, rule : str) -> None:
        field, ranges = rule.split(":")
        ranges = findall(r"\d+-\d+", ranges)
        ranges = [r.split("-") for r in ranges]
        self.rules[field] = [[int(v) for v in r] for r in ranges]

    def addRules(self, rules : List[str]) -> None:
        for rule in rules:
            self.addRule(rule)

    def validate_all(self, values : List[int]) -> dict:
        result = {
            "invalid_values": [],
        }
        for v in values:
            is_valid = False
            for field_ranges in self.rules.values():
                for min_val, max_val in field_ranges:
                    if v >= min_val and v <= max_val:
                        is_valid = True
                        break
                if is_valid == True:
                    break
            if not is_valid:
                result["invalid_values"].append(v)
        return result

if __name__ == "__main__":
    dir_path = dirname(realpath(__file__))
    with open(dir_path + "/input.txt") as input_file:
        ignore = [
            "\n",
            "your ticket:\n",
            "nearby tickets:\n",
        ]
        ignored = 0
        sum_of_invalid = 0
        validator = TicketValidator()
        for line in input_file:
            if line in ignore:
                ignored += 1
                continue
            if ignored == 0:
                validator.addRule(line)
            elif ignored > 0:
                result = validator.validate_all([int(v) for v in line.split(",")])
                sum_of_invalid += sum(result["invalid_values"])

        print(sum_of_invalid)


# import third-party libraries
from tabulate import tabulate
from colorama import init as colorama_init, Fore, Style
import httpx, asyncio, aiofiles

# import Python's standard libraries
from typing import Iterable, Literal, Union, Optional, Generator
from pathlib import Path
from time import perf_counter_ns
from json import JSONDecodeError

# import local files
if (__package__ is None or __package__ == ""):
    from constants import CONSTANTS
    from spinner import Spinner
    from functional import get_input, print_success, print_warning, print_danger, load_configs
    from user_data import load_danbooru_auth, save_danbooru_auth
    from errors import TooManyRequests, ApiKeyError
else:
    from .constants import CONSTANTS
    from .spinner import Spinner
    from .functional import get_input, print_success, print_warning, print_danger, load_configs
    from .user_data import load_danbooru_auth, save_danbooru_auth
    from .errors import TooManyRequests, ApiKeyError

# -------- Danbooru Data Objects --------
class Post():
    def __init__(self, id:int, data:dict, detailed:bool = False) -> None:
        if "bad_id" in data.get("tag_string_meta"):
            self.source = None
        elif data.get("pixiv_id") is not None:
            self.source = f"https://www.pixiv.net/en/artworks/{data.get('pixiv_id')}" # If pixiv source
        else:
            self.source = data.get("source")
            if self.source == "":
                self.source = None

        self.id = id
        self.md5 = data.get("md5")
        self.tag_string_general = set(data.get("tag_string_general").split())
        self.is_banned = data.get("is_banned")

        if detailed:
            characters = data.get("tag_string_character")
            if characters == "":
                self.characters = f"{Fore.LIGHTYELLOW_EX}N.A.{Style.RESET_ALL}"
            else:
                characters = characters.split()
                self.characters = [character.replace("_", " ").title() for character in characters[:3]]
                if len(characters) > 3:
                    self.characters.append(f"{Fore.LIGHTYELLOW_EX}And {len(characters) - 3} more...{Style.RESET_ALL}")
                self.characters = "\n".join(self.characters)

            artists = data.get("tag_string_artist")
            if artists == "":
                self.artists = f"{Fore.LIGHTYELLOW_EX}N.A.{Style.RESET_ALL}"
            else:
                artists = artists.split()
                self.artists = [artist.replace("_", " ").title() for artist in artists[:3]]
                if len(artists) > 3:
                    self.artists.append(f"{Fore.LIGHTYELLOW_EX}And {len(artists) - 3} more...{Style.RESET_ALL}")
                self.artists = "\n".join(self.artists)

            copyrights = data.get("tag_string_copyright")
            if copyrights == "":
                self.copyrights = f"{Fore.LIGHTYELLOW_EX}N.A.{Style.RESET_ALL}"
            else:
                copyrights = copyrights.split()
                self.copyrights = [copyright.replace("_", " ").title() for copyright in copyrights[:3]]
                if len(copyrights) > 3:
                    self.copyrights.append(f"{Fore.LIGHTYELLOW_EX}And {len(copyrights) - 3} more...{Style.RESET_ALL}")
                self.copyrights = "\n".join(self.copyrights)

        if image_size == "small": # This makes sense don't think too much about it.
            self.file_url = data.get("large_file_url")
        elif image_size == "large": # Both have the same value in the event the image is small.
            self.file_url = data.get("file_url")

    def __str__(self):
        try:
            return tabulate(
                tabular_data = [(self.id, self.artists, self.characters, self.copyrights)],
                headers = ("Post", "Artist(s)", "Character(s)", "Copyright(s)"),
                tablefmt = "simple_grid",
                disable_numparse = True,
            )
        except AttributeError:
            raise AttributeError("\"detailed\" variable must be set to \"True\"!")

    def check_valid(self) -> Optional[str]:
        source = ""
        if self.is_banned:
            if self.source is not None:
                source = f"The image can be downloaded manually from '{self.source}'."
                message = f"{Fore.YELLOW}Post {self.id} has been banned from Danbooru. {source}{Style.RESET_ALL}"
        elif self.tag_string_general.intersection({"loli", "shota", "toddlercon"}):
            if self.source is not None:
                source = f"The image can be downloaded manually from '{self.source}'."
            message = f"{Fore.YELLOW}Post {self.id} is Danbooru Gold exclusive only. {source}{Style.RESET_ALL}"
        else:
            message = None
        return message

    async def save(self) -> None:
        message = self.check_valid()
        if message is not None:
            print_warning(message)
            return

        filepath = directory.joinpath(f"Post_{self.id}.{self.file_url.rsplit('.', 1)[-1]}")
        if filepath.is_file():
            print_warning(f"File \"{filepath.name}\" already exists.")

        async with httpx.AsyncClient() as client:
            await download_async_group(self.file_url, client, filepath)

class Pool():
    def __init__(self, id:int, data:dict) -> None:
        self.id = id
        self.name = data.get("name").replace("_", " ").title()
        self.posts = data.get("post_ids")
        self.post_count = len(self.posts)

    def __str__(self) -> str:
        return tabulate(
            tabular_data = [(self.id, self.post_count, self.name)],
            headers = ("Pool", "Posts", "Name"),
            tablefmt = "simple_grid",
            disable_numparse = True,
            maxcolwidths = [None, None, 75]
        )

    async def save(self, mode:str, limit:int, direction:str, pivot:int) -> None:
        if mode == "id":
            if direction == ">":
                posts = sorted(filter(lambda id: id > pivot, self.posts))[:limit]
            elif direction == "<":
                posts = sorted(filter(lambda id: id < pivot, self.posts))[-limit:]

        elif mode == "order":
            if direction == ">":
                posts = self.posts[pivot: limit + pivot]
            elif direction == "<":
                posts = self.posts[max(0, pivot - limit): pivot]

        now = perf_counter_ns()
        complete = 0
        spinner_message = f"Downloading images.. {{complete}}/{len(posts)} ({{percent}}%)"

        folder = directory.joinpath(self.name)
        if folder.is_dir():
            print_warning(f"Folder \"{folder.name}\" already exists.")
        folder.mkdir(parents=True, exist_ok=True)

        messages = []
        with Spinner(spinner_message.format(complete=0, percent=0), colour="yellow", spinner_position="left", spinner_type="material", completion_msg="Done", cancelled_msg="Cancelled\n") as spinner:
            async with httpx.AsyncClient() as client:
                try:
                    for async_group in generate_async_group(posts):
                        data = await asyncio.gather(*(fetch_async_group(
                                    client = client,
                                    post_id = post)
                                for post in async_group))

                        valid = []
                        for index, id in enumerate(async_group):
                            response = data[index]
                            message = verify_response(id, "post", response)
                            if message is not None:
                                messages.append(message)
                            else:
                                post = Post(id, response.json())
                                message = post.check_valid()
                                messages.append(message) if message is not None else valid.append(post)

                        if mode == "id": # Post_<id>
                            await asyncio.gather(*(download_async_group(
                                    url = post.file_url,
                                    client = client,
                                    filepath = folder.joinpath(f"Post_{post.id}.{post.file_url.rsplit('.', 1)[-1]}"))
                                for post in valid))

                        elif mode == "order": # Page_<num>
                            await asyncio.gather(*(download_async_group(
                                    url = post.file_url,
                                    client = client,
                                    filepath = folder.joinpath(f"Page_{complete+index+1}.{post.file_url.rsplit('.', 1)[-1]}"))
                                for index, post in enumerate(valid)))

                        complete += len(async_group)
                        spinner.message = spinner_message.format(complete=complete, percent=round(complete/self.post_count*100, 2))

                except TooManyRequests as error:
                    spinner.stop(error=True)
                    print_danger(error)
                except KeyboardInterrupt:
                    spinner.stop(manually_stopped=True)

        print(f"\n[Process complete in {round((perf_counter_ns() - now)/1000000000, 2)}s]")
        if messages:
            if get_input("Some images were not downloaded. See more information? (Y/n):", inputs=("y", "n"), default="y") == "y":
                print(*messages, sep="\n")

class Artist():
    def __init__(self, id:int, data:dict):
        self.id = id
        self.name = data.get("name")

        self.group = data.get("group_name").replace("_", " ").title()
        if self.group == "":
            self.group = f"{Fore.LIGHTYELLOW_EX}N.A.{Style.RESET_ALL}"

        aliases = data.get("other_names")
        self.alias = [alias.replace("_", " ").title() for alias in aliases[:3]]
        if len(aliases) > 3:
            self.alias.append(f"{Fore.LIGHTYELLOW_EX}And {len(aliases) - 3} more...{Style.RESET_ALL}")
        self.alias = "\n".join(self.alias)

        self.is_banned = data.get("banned")

    def __str__(self):
        return tabulate(
            tabular_data=[(self.id, self.name.replace("_", " ").title(), self.group, self.alias)],
            headers = ("ID", "Name", "Group", "Alias"),
            tablefmt = "simple_grid",
            disable_numparse = True
        )

    async def save(self, limit:int, direction:str, pivot:int) -> None:
        response = httpx.get(
            url = "https://danbooru.donmai.us/posts.json", auth = get_auth(),
            params = {
                "limit": limit,
                "tags": self.name,
                "page": ("a" if direction == ">" else "b") + str(pivot) # b<id>, a<id>
            })

        message = verify_response(self.id, "artist", response)
        if message is not None:
            print(message)
            return

        folder = directory.joinpath(self.name)
        if folder.is_dir():
            print_warning(f"Folder \"{folder.name}\" already exists.")
        folder.mkdir(parents=True, exist_ok=True)

        messages = {"gold": 0, "banned": 0, "source":[]}
        now = perf_counter_ns()
        complete = 0
        spinner_message = f"Downloading images.. {{complete}}/{len(response.json())} ({{percent}}%)"

        with Spinner(spinner_message.format(complete=0, percent=0), colour="yellow", spinner_position="left", spinner_type="material", completion_msg="Done", cancelled_msg="Cancelled\n") as spinner:
            async with httpx.AsyncClient() as client:
                # try:
                    for async_group in generate_async_group(response.json()):
                        valid = []
                        for data in async_group:
                            if data.get("id") is None:
                                if data.get("pixiv_id") is not None:
                                    messages.get("source").append(f"https://www.pixiv.net/en/artworks/{data.get('pixiv_id')}") # If pixiv source
                                elif "bad_id" not in data.get("tag_string_meta"):
                                    messages.get("source").append(data.get("source"))

                                if self.is_banned:
                                    messages["banned"] += 1
                                elif set(data.get("tag_string_general")).intersection({"loli", "shota", "toddlercon"}):
                                    messages["gold"] += 1
                            else:
                                valid.append(Post(data.get("id"), data))

                            await asyncio.gather(*(download_async_group(
                                url = post.file_url,
                                client = client,
                                filepath = folder.joinpath(f"Post_{post.id}.{post.file_url.rsplit('.', 1)[-1]}"))
                            for post in valid))

                        complete += len(async_group)
                        spinner.message = spinner_message.format(complete=complete, percent=round(complete/limit*100, 2))

                # except TooManyRequests as error:
                #     spinner.stop(error=True)
                #     print_danger(error)
                #     for task in asyncio.all_tasks():
                #         task.cancel()
                # except KeyboardInterrupt:
                #     spinner.stop(manually_stopped=True)
                #     for task in asyncio.all_tasks():
                #         task.cancel()
                #     print("Cancelling")

        print(f"\n[Process complete in {round((perf_counter_ns() - now)/1000000000, 2)}s]")
        if messages.get("gold") + messages.get("banned") != 0:
            sources = '\n- '.join(messages.get('source'))
            if get_input("Some images were not downloaded. See more information? (Y/n):", inputs=("y", "n"), default="y") == "y":
                print(f"""\
Danbooru Gold Exclusive: {Fore.LIGHTYELLOW_EX}{messages.get('gold')}{Style.RESET_ALL}\t\
Banned: {Fore.LIGHTYELLOW_EX}{messages.get('banned')}{Style.RESET_ALL}
The following images can be manually downloaded:
-{sources}""")

# ---------- Helper Functions ----------
def fetch_resource(id:int, type:str) -> httpx.Response:
    if type not in {"post", "pool", "artist"}:
        raise Exception("Type variable can only be \"post\", \"pool\" or \"artist\"!")

    response = httpx.get(url = f"https://danbooru.donmai.us/{type}s/{id}.json", auth = get_auth())
    return response

def verify_response(id:int, type:str, response:httpx.Response) -> Optional[str]:
    if response.status_code > 399:
        if response.status_code == 401:
            message = f"{Fore.LIGHTRED_EX}Error with API key/username, please re-configure.{Style.RESET_ALL}"
        elif response.status_code == 403:
            message = f"{Fore.LIGHTRED_EX}API key does not have the necessary permissions, please re-configure.{Style.RESET_ALL}"
        if response.status_code == 404:
            message = f"{Fore.LIGHTRED_EX}{type.capitalize()} {id} was not found.{Style.RESET_ALL}"
        else:
            try:
                data = response.json()
                message = f"{Fore.LIGHTRED_EX} Error {response.status_code}: {data.get('message')} [{data.get('error')}]{Style.RESET_ALL}"
            except JSONDecodeError:
                message = f"{Fore.LIGHTRED_EX}Error {response.status_code}: Unknown error{Style.RESET_ALL}"
    else:
        message = None
    return message

def generate_async_group(post_ids:Iterable[int]) -> Generator[int, None, None]:
    start, length = 0, len(post_ids)

    for limit, step in group_generation:
        end = (length + 1) if limit is None else (start + limit)

        for index in range(start, end, step):
                yield post_ids[index: min(end, index+step)]

        if limit is None or length <= limit:
            return
        start += limit

def legacy_generate_async_group(post_ids:Iterable[int], length:int=3) -> Generator[int, None, None]:
    if length > 5:
        print(f"{Fore.YELLOW}Warning: An async group of larger than 5 (e.g. 10) may result in '429: Too Many Requests' for large pools.{Fore.RESET}")
    for index in range(0, len(post_ids), length):
        yield post_ids[index: index+length]

async def fetch_async_group(client:httpx.AsyncClient, post_id:int) -> dict:
    response = await client.get(f"https://danbooru.donmai.us/posts/{post_id}.json")
    if response.status_code == 429:
        raise TooManyRequests("The server is unable to handle the large number of posts queries; please try again later.")
    return response

async def download_async_group(url:Union[str, Path], client:httpx.AsyncClient, filepath:Union[str, Path]) -> None:
    async with client.stream("GET", url) as response, aiofiles.open(filepath, "wb") as file:
        async for chunk in response.aiter_bytes(CONSTANTS.CHUNK_SIZE):
            await file.write(chunk)

def get_auth():
    if username is None or api_key is None:
        return None
    else:
        return (username, api_key)

# --------- User Data Management ---------
class UserProfile():
    # Default values from: https://danbooru.donmai.us/profile.json
    def __init__(self, id:Optional[int]=None, name:str="Anonymous", level:int=0, level_string:str="Anonymous"):
        self.id = id
        self.name = name
        self.level = level
        self.level_string = level_string

    def __str__(self):
        username = "Guest" if self.name == "Anonymous" else self.name
        user_level = "Not logged in" if self.level_string == "Anonymous" else self.level_string
        return f"Danbooru: {username} ({user_level})"

    def __repr__(self):
        return tabulate(
            tabular_data = [(self.id, self.name, f"{self.level_string} [{self.level}]")],
            headers = ("User ID", "Username", "Level"),
            tablefmt = "simple_grid",
            disable_numparse = True
        )

def get_user_data(username:Optional[str]=None, api_key:Optional[str]=None) -> UserProfile:
    if username is None or api_key is None:
        auth = get_auth()
    else:
        auth = (username, api_key)

    response = httpx.get(url = "https://danbooru.donmai.us/profile.json", auth = auth)
    message = verify_response(username, "User", response)

    if message is not None:
        raise ValueError(message)
    else:
        data = response.json()
        return UserProfile(data.get("id"), data.get("name"), data.get("level"), data.get("level_string"))

# ------ Initialise global values ------
username:Optional[str] = None
api_key:Optional[str] = None
image_size:Literal["large", "small"] = "large"
directory:Path = Path(load_configs().download_directory)
group_generation:tuple[tuple[Optional[int], int]] = ((50, 10), (50, 5), (400, 4), (None, 3))
user_profile:UserProfile = UserProfile(None, "Anonymous", 0, "Anonymous") # Default values for non-users

def auto_init():
    global username, api_key, user_profile

    user_data = load_danbooru_auth()
    if user_data is not None:
        username = user_data.get("username")
        api_key = user_data.get("api_key")
    else: # Something wrong, use default configurations
        username, api_key = None, None

    try:
        user_profile = get_user_data()
    except ValueError as error:
        # print_danger(error)
        user_profile = UserProfile(None, "Anonymous", 0, "Anonymous")

colorama_init()

__all__ = ["api_key", "username", "user_profile", "directory", "image_size", "fetch_resource",
        "generate_async_group", "fetch_async_group", "download_async_group",
        "Post", "Pool", "Artist"]

if __name__ == "__main__":
    # print(load_danbooru_auth())
    print(user_profile)
    print(repr(user_profile))