import asyncio
import logging
import sys
import time

import pyautogui
import pydirectinput
import qasync
from PySide6 import QtWidgets, QtCore
from front import Ui_MainWindow

def use_quick_access_inventory():
    print("Use quick access of inventory")
    pydirectinput.keyDown('1')
    time.sleep(0.1)
    pydirectinput.keyUp('1')
    pydirectinput.keyDown('2')
    time.sleep(0.1)
    pydirectinput.keyUp('2')
    pydirectinput.keyDown('3')
    time.sleep(0.1)
    pydirectinput.keyUp('3')
    pydirectinput.keyDown('4')
    time.sleep(0.1)
    pydirectinput.keyUp('4')

def logout():
    print("Logout")
    time.sleep(0.1)
    pydirectinput.leftClick(1264, 967)
    time.sleep(1)
    pydirectinput.leftClick(655, 573)

@qasync.asyncSlot()
async def terminate_spot(time_attack):
    pydirectinput.keyDown('1')
    await asyncio.sleep(0.1)
    pydirectinput.keyUp('1')
    await asyncio.sleep(0.1)
    pydirectinput.keyDown('space')
    await asyncio.sleep(time_attack)
    pydirectinput.keyUp('space')
    pydirectinput.keyDown('z')
    await asyncio.sleep(0.1)
    pydirectinput.keyUp('z')

@qasync.asyncSlot()
async def change_chanel(chanel, time_change):

    if chanel < 0 or chanel > 4:
        print(f'Chanel cant be {chanel}')
        return

    get_window("Insomnia Helper", False)
    get_window('InsomniaMT2', True)
    print(f'Chanel is changing to chanel {chanel}')

    if chanel == 1:
        await asyncio.sleep(0.1)
        pydirectinput.leftClick(1156, 81)
        await asyncio.sleep(time_change)
        return

    if chanel == 2:
        await asyncio.sleep(0.1)
        pydirectinput.leftClick(1154, 102)
        await asyncio.sleep(time_change)
        return

    if chanel == 3:
        await asyncio.sleep(0.1)
        pydirectinput.leftClick(1160, 127)
        await asyncio.sleep(time_change)
        return

    if chanel == 4:
        await asyncio.sleep(0.1)
        pydirectinput.leftClick(1174, 147)
        await asyncio.sleep(time_change)
        return

def use_main_inventory_slots(number):
    if number < 1 or number > 5:
        return False

    print(f'Use slot {number} of main inventory')

    if number == 1:
        time.sleep(0.1)
        pydirectinput.leftClick(1137, 645)
        return True

    if number == 2:
        time.sleep(0.1)
        pydirectinput.leftClick(1168, 645)
        return True

    if number == 3:
        time.sleep(0.1)
        pydirectinput.leftClick(1200, 645)
        return True

    if number == 4:
        time.sleep(0.1)
        pydirectinput.leftClick(1234, 645)
        return True

    if number == 5:
        time.sleep(0.1)
        pydirectinput.leftClick(1264, 645)
        return True

def open_main_inventory():
    print("Opening main inventory")
    time.sleep(0.1)
    pydirectinput.leftClick(1192, 970)

def get_window(title, client):
    while True:
        hwnd = pyautogui.getWindowsWithTitle(title)

        if len(hwnd) <= 0:
            return None

        if len(hwnd) >= 2:
            print(f'Number of finding windows is {len(hwnd)}')
            window_to_remove = pyautogui.getWindowsWithTitle("InsomniaMT2 Klient")
            window_to_remove[0].maximize()
            decison = pyautogui.confirm(text='Remove this window?', title='Remove', buttons=['OK', 'Cancel'])
            if decison == 'Cancel':
                return None

            if decison == 'OK':
                window_to_remove[0].close()
                time.sleep(0.1)

        if len(hwnd) == 1:
            hwnd[0].activate()
            if client:
                hwnd[0].moveTo(0, 0)
            else:
                hwnd[0].moveTo(1280, 0)
            return hwnd

class MainWindow(QtWidgets.QMainWindow, Ui_MainWindow):
    def __init__(self):
        super(MainWindow, self).__init__()
        self.setupUi(self)
        self.setWindowTitle("Insomnia Helper")
        self.stackedWidget.setCurrentIndex(0)

        # Variable
        self.can_running = False
        self.counter_chanel = 1

        # Connect push button
        self.autoDropping.clicked.connect(self.auto_dropping)
        self.bossHelpper.clicked.connect(self.boss_helper)
        self.back_button.clicked.connect(self.back_start_page)
        self.back_button_1.clicked.connect(self.back_start_page)
        self.start_dropping.clicked.connect(self.start_dropping_fun)

    def auto_dropping(self):
        self.stackedWidget.setCurrentIndex(1)

    def boss_helper(self):
        self.stackedWidget.setCurrentIndex(2)

    def back_start_page(self):
        self.stackedWidget.setCurrentIndex(0)

    def start_dropping_fun(self):
        print("button clicked")
        self.can_running = not self.can_running

        if self.can_running:
            self.main()

    def keyPressEvent(self, event):
        key = event.key()

        # if key == 61:
        #     self.can_running = True
        #     print("Key started is clicked")
        #     self.main()
        #     return
        #
        # if key == 45:
        #     print("Key stopped is clicked")
        #     self.can_running = False
        #     return

        print(f"Found key {key} --> dont have action for this key")

    @qasync.asyncSlot()
    async def main(self):
        print("Dropping is started")
        try:
            while True:
                print(self.can_running)
                if not self.can_running:
                    print("Dropping is stopped")
                    break
                await self.auto_dropping_fun()
        except:
            print("Error with dropping, try again")

    @qasync.asyncSlot()
    async def auto_dropping_fun(self):

        if not self.ch1_check_box.isChecked() and self.counter_chanel == 1:
            self.counter_chanel += 1

        if not self.ch2_check_box.isChecked() and self.counter_chanel == 2:
            self.counter_chanel += 1

        if not self.ch3_check_box.isChecked() and self.counter_chanel == 3:
            self.counter_chanel += 1

        if not self.ch4_check_box.isChecked() and self.counter_chanel == 4:
            self.counter_chanel += 1

        if self.counter_chanel > 4:
            self.counter_chanel = 1
            await asyncio.sleep(0.1)
            return

        await terminate_spot(self.atact_time.value())
        await change_chanel(self.counter_chanel, self.change_chanel_time.value())
        self.counter_chanel += 1

def start_application():  # Start Application with qasync
    app = QtWidgets.QApplication(sys.argv)
    loop = qasync.QEventLoop(app)
    window = MainWindow()
    window.show()
    logging.info("Starting application Insomnia bot ...")
    with loop:
        loop.run_forever()

start_application()


# -*- coding: utf-8 -*-

from django.shortcuts import render, redirect, get_object_or_404
from django.urls import reverse
from django.contrib import messages
from django.utils.translation import gettext as _

from djconfig import config

from spirit.core.utils.views import is_post, post_data
from spirit.core.utils.paginator import yt_paginate
from spirit.core.utils.decorators import administrator_required
from .forms import CommentFlagForm
from ..models import CommentFlag, Flag

@administrator_required
def detail(request, pk):
    flag = get_object_or_404(CommentFlag, pk=pk)
    form = CommentFlagForm(
        user=request.user,
        data=post_data(request),
        instance=flag)

    if is_post(request) and form.is_valid():
        form.save()
        messages.info(request, _("The flag has been moderated!"))
        return redirect(reverse("spirit:admin:flag:index"))

    flags = yt_paginate(
        Flag.objects.filter(comment=flag.comment),
        per_page=config.comments_per_page,
        page_number=request.GET.get('page', 1)
    )

    return render(
        request=request,
        template_name='spirit/comment/flag/admin/detail.html',
        context={
            'flag': flag,
            'flags': flags,
            'form': form})

@administrator_required
def _index(request, queryset, template):
    flags = yt_paginate(
        queryset,
        per_page=config.comments_per_page,
        page_number=request.GET.get('page', 1)
    )
    context = {'flags': flags, }
    return render(request, template, context)

def opened(request):
    return _index(
        request,
        queryset=CommentFlag.objects.filter(is_closed=False),
        template='spirit/comment/flag/admin/open.html'
    )

def closed(request):
    return _index(
        request,
        queryset=CommentFlag.objects.filter(is_closed=True),
        template='spirit/comment/flag/admin/closed.html'
    )


import pandas as pd
import math
from sklearn import linear_model
import numpy as np

def predict_using_sklearn():
    test_scores = pd.read_csv(r'C:\Users\Polina\Desktop\Python\Pandas\test_scores.csv')
    reg = linear_model.LinearRegression()
    reg.fit(test_scores[['math']], test_scores.cs)
    return reg.coef_, reg.intercept_

def gradient_descent(x,y):
    m_curr=b_curr=0
    iterations = 1000000
    n = len(x)
    learning_rate = 0.0002
    cost_previous = 0
    for i in range(iterations):
        y_predicted = m_curr*x + b_curr
        cost = (1/n)*sum([val**2 for val in (y-y_predicted)])
        md = -(2/n)*sum(x*(y-y_predicted))
        bd = -(2/n)*sum(y -y_predicted)
        m_curr = m_curr- learning_rate*md
        b_curr = b_curr - learning_rate*bd
        if math.isclose(cost, cost_previous, rel_tol=1e-20):
            break
        cost_previous = cost
        return m_curr, b_curr

if __name__ == '__main__':

    df = pd.read_csv(r"C:\Users\Polina\Desktop\Python\Pandas\test_scores.csv")
    x = np.array(df.math)
    y = np.array(df.cs)

    m, b = gradient_descent(x,y)
    print("Using gradient descent function: Coef {} Intercept {}".format(m, b))

    m_sklearn, b_sklearn = predict_using_sklearn()
    print("Using sklearn: Coef {} Intercept {}".format(m_sklearn,b_sklearn))


from syscore.objects import arg_not_supplied
from syscore.genutils import flatten_list
from dataclasses import dataclass
import pandas as pd

EMPTY_INSTRUMENT = ""

class futuresInstrument(object):
    def __init__(self, instrument_code: str):
        self._instrument_code = instrument_code

    @property
    def instrument_code(self):
        return self._instrument_code

    def empty(self):
        return self.instrument_code == EMPTY_INSTRUMENT

    @classmethod
    def create_from_dict(futuresInstrument, input_dict):
        # Might seem pointless, but (a) is used in original code, (b) gives a nice consistent feel
        return futuresInstrument(input_dict["instrument_code"])

    def as_dict(self):
        # Might seem pointless, but (a) is used in original code, (b) gives a nice consistent feel
        return dict(instrument_code=self.instrument_code)

    def __eq__(self, other):
        return self.instrument_code == other.instrument_code

    @property
    def key(self):
        return self.instrument_code

    def __repr__(self):
        return str(self.instrument_code)

@dataclass
class instrumentMetaData:
    Description: str = ""
    Pointsize: float = 0.0
    Currency: str = ""
    AssetClass: str = ""
    Slippage: float = 0.0
    PerBlock: float = 0.0
    Percentage: float = 0.0
    PerTrade: float = 0.0

    def as_dict(self) -> dict:
        keys = self.__dataclass_fields__.keys()
        self_as_dict = dict([(key, getattr(self, key)) for key in keys])

        return self_as_dict

    @classmethod
    def from_dict(instrumentMetaData, input_dict):
        keys = instrumentMetaData.__dataclass_fields__.keys()
        args_list = [input_dict[key] for key in keys]

        return instrumentMetaData(*args_list)

@dataclass
class futuresInstrumentWithMetaData:
    instrument: futuresInstrument
    meta_data: instrumentMetaData

    @property
    def instrument_code(self) -> str:
        return self.instrument.instrument_code

    @property
    def key(self) -> str:
        return self.instrument_code

    def as_dict(self) -> dict:
        meta_data_dict = self.meta_data.as_dict()
        meta_data_dict["instrument_code"] = self.instrument_code

        return meta_data_dict

    @classmethod
    def from_dict(futuresInstrumentWithMetaData, input_dict):
        instrument_code = input_dict.pop("instrument_code")
        instrument = futuresInstrument(instrument_code)
        meta_data = instrumentMetaData.from_dict(input_dict)

        return futuresInstrumentWithMetaData(instrument, meta_data)

    @classmethod
    def create_empty(futuresInstrumentWithMetaData):
        instrument = futuresInstrument(EMPTY_INSTRUMENT)
        meta_data = instrumentMetaData()

        instrument_with_metadata = futuresInstrumentWithMetaData(instrument, meta_data)

        return instrument_with_metadata

    def empty(self):
        return self.instrument.empty()

class listOfFuturesInstrumentWithMetaData(list):
    def as_df(self):
        instrument_codes = [
            instrument_object.instrument_code for instrument_object in self
        ]
        meta_data_keys = [
            instrument_object.meta_data.as_dict().keys() for instrument_object in self
        ]
        meta_data_keys_flattened = flatten_list(meta_data_keys)
        meta_data_keys_unique = list(set(meta_data_keys_flattened))

        meta_data_as_lists = dict(
            [
                (
                    metadata_name,
                    [
                        getattr(instrument_object.meta_data, metadata_name)
                        for instrument_object in self
                    ],
                )
                for metadata_name in meta_data_keys_unique
            ]
        )

        meta_data_as_dataframe = pd.DataFrame(
            meta_data_as_lists, index=instrument_codes
        )

        return meta_data_as_dataframe

class assetClassesAndInstruments(dict):
    def __repr__(self):
        return str(self.as_pd())

    def get_instrument_list(self) -> list:
        return list(self.keys())

    @classmethod
    def from_pd_series(self, pd_series: pd.Series):
        instruments = list(pd_series.index)
        asset_classes = list(pd_series.values)
        as_dict = dict(
            [
                (instrument_code, asset_class)
                for instrument_code, asset_class in zip(instruments, asset_classes)
            ]
        )

        return assetClassesAndInstruments(as_dict)

    def all_asset_classes(self) -> list:
        asset_classes = list(self.values())
        unique_asset_classes = list(set(asset_classes))
        unique_asset_classes.sort()

        return unique_asset_classes

    def as_pd(self) -> pd.Series:
        instruments = [key for key in self.keys()]
        asset_classes = [value for value in self.values()]

        return pd.Series(asset_classes, index=instruments)

    def all_instruments_in_asset_class(
        self, asset_class: str, must_be_in=arg_not_supplied
    ) -> list:

        asset_class_instrument_list = [
            instrument
            for instrument, item_asset_class in self.items()
            if item_asset_class == asset_class
        ]

        if must_be_in is arg_not_supplied:
            return asset_class_instrument_list

        ## we need to filter
        filtered_asset_class_instrument_list = [
            instrument
            for instrument in asset_class_instrument_list
            if instrument in must_be_in
        ]

        return filtered_asset_class_instrument_list

class instrumentCosts(object):
    def __init__(
        self,
        price_slippage: float = 0.0,
        value_of_block_commission: float = 0.0,
        percentage_cost: float = 0.0,
        value_of_pertrade_commission: float = 0.0,
    ):
        self._price_slippage = price_slippage
        self._value_of_block_commission = value_of_block_commission
        self._percentage_cost = percentage_cost
        self._value_of_pertrade_commission = value_of_pertrade_commission

    @classmethod
    def from_meta_data(instrumentCosts, meta_data: instrumentMetaData):
        return instrumentCosts(
            price_slippage=meta_data.Slippage,
            value_of_block_commission=meta_data.PerBlock,
            percentage_cost=meta_data.Percentage,
            value_of_pertrade_commission=meta_data.PerTrade,
        )

    def __repr__(self):
        return (
            "instrumentCosts slippage %f block_commission %f percentage cost %f per trade commission %f "
            % (
                self.price_slippage,
                self.value_of_block_commission,
                self.percentage_cost,
                self.value_of_pertrade_commission,
            )
        )

    @property
    def price_slippage(self):
        return self._price_slippage

    @property
    def value_of_block_commission(self):
        return self._value_of_block_commission

    @property
    def percentage_cost(self):
        return self._percentage_cost

    @property
    def value_of_pertrade_commission(self):
        return self._value_of_pertrade_commission

    def calculate_cost_percentage_terms(
        self, blocks_traded: float, block_price_multiplier: float, price: float
    ) -> float:
        cost_in_currency_terms = self.calculate_cost_instrument_currency(
            blocks_traded, block_price_multiplier=block_price_multiplier, price=price
        )

        value_per_block = price * block_price_multiplier
        total_value = blocks_traded * value_per_block
        cost_in_percentage_terms = cost_in_currency_terms / total_value

        return cost_in_percentage_terms

    def calculate_cost_instrument_currency(
        self, blocks_traded: float, block_price_multiplier: float, price: float
    ) -> float:

        value_per_block = price * block_price_multiplier
        slippage = self.calculate_slippage_instrument_currency(
            blocks_traded, block_price_multiplier=block_price_multiplier
        )

        commission = self.calculate_total_commission(
            blocks_traded, value_per_block=value_per_block
        )

        return slippage + commission

    def calculate_total_commission(self, blocks_traded: float, value_per_block: float):
        ### YOU WILL NEED TO CHANGE THIS IF YOUR BROKER HAS A MORE COMPLEX STRUCTURE
        per_trade_commission = self.calculate_per_trade_commission()
        per_block_commission = self.calculate_cost_per_block_commission(blocks_traded)
        percentage_commission = self.calculate_percentage_commission(
            blocks_traded, value_per_block
        )

        return max([per_block_commission, per_trade_commission, percentage_commission])

    def calculate_slippage_instrument_currency(
        self, blocks_traded: float, block_price_multiplier: float
    ) -> float:
        return abs(blocks_traded) * self.price_slippage * block_price_multiplier

    def calculate_per_trade_commission(self):
        return self.value_of_pertrade_commission

    def calculate_cost_per_block_commission(self, blocks_traded):
        return abs(blocks_traded) * self.value_of_block_commission

    def calculate_percentage_commission(self, blocks_traded, price_per_block):
        trade_value = self.calculate_trade_value(blocks_traded, price_per_block)
        return self._percentage_cost * trade_value

    def calculate_trade_value(self, blocks_traded, value_per_block):
        return abs(blocks_traded) * value_per_block


from __future__ import annotations

import pytest

@pytest.mark.parametrize(
    "variables, expected_data",
    [
        (
            {"name": r"pants_explorer\."},
            {
                "rules": [
                    {"name": "pants_explorer.server.graphql.rules.get_graphql_uvicorn_setup"},
                ]
            },
        ),
        (
            {"name": r"\.graphql\."},
            {
                "rules": [
                    {"name": "pants_explorer.server.graphql.rules.get_graphql_uvicorn_setup"},
                ]
            },
        ),
        (
            {"limit": 0},
            {"rules": []},
        ),
        (
            {"limit": 0},
            {"rules": []},
        ),
    ],
)
def test_rules_query(
    schema, queries: str, variables: dict, expected_data: dict, context: dict
) -> None:
    actual_result = schema.execute_sync(
        queries, variable_values=variables, context_value=context, operation_name="TestRulesQuery"
    )
    assert actual_result.errors is None
    assert actual_result.data == expected_data


from django.db import models
from django.db.models import Case, Count, IntegerField, When

class CountryManager(models.Manager):
    def aggregate_integration_statuses(self):
        from proco.connection_statistics.models import CountryWeeklyStatus
        return self.get_queryset().aggregate(
            countries_joined=Count(Case(When(
                last_weekly_status__integration_status__in=[
                    CountryWeeklyStatus.SCHOOL_MAPPED,
                    CountryWeeklyStatus.STATIC_MAPPED,
                    CountryWeeklyStatus.REALTIME_MAPPED,
                ], then=1),
                output_field=IntegerField())),
            countries_connected_to_realtime=Count(Case(When(
                last_weekly_status__integration_status=CountryWeeklyStatus.REALTIME_MAPPED, then=1),
                output_field=IntegerField())),
            countries_with_static_data=Count(Case(When(
                last_weekly_status__integration_status=CountryWeeklyStatus.STATIC_MAPPED, then=1),
                output_field=IntegerField()),
            ),
            countries_with_static_and_realtime_data=Count(Case(When(
                last_weekly_status__integration_status= (CountryWeeklyStatus.STATIC_MAPPED or CountryWeeklyStatus.REALTIME_MAPPED), then=1),
                output_field=IntegerField()) ,
            ),
        )


import json
from statistics import mean
import numpy as np
import os
from bokeh.plotting import output_file, figure, save
from bokeh.layouts import gridplot
from src.utils.tools import listdir, hash_append

def combined(ids, name, legend=None, y_range=(0, 900)):
    summaries = {}
    episodes = []
    for key_ in settings:  # [key_ for key_ in settings if len(settings[key_]) > 1]:
        id = [f.split(' := ') for f in key_.split('\n') if f.split(' := ')[0] == 'id'][0][1]
        if not any([id == id_ for id_ in ids]):
            continue
        rewards = {}
        for res_folder in settings[key_]:
            with open(os.path.join(res_folder, 'metrics.json'), 'r') as f:
                metrics = json.load(f)
            for episode, score in zip(metrics['episodes'], metrics['t_scores']):
                if score is not None and episode % 25 == 0:
                    hash_append(rewards, episode, score)
        episodes = sorted(rewards.keys())
        episodes = episodes[:min(39, len(episodes))]
        quart1 = [np.percentile(rewards[ep], 25) for ep in episodes]
        median = [np.percentile(rewards[ep], 50) for ep in episodes]
        quart3 = [np.percentile(rewards[ep], 75) for ep in episodes]
        summaries[id] = (quart1, median, quart3)
    COLORS = ("royalblue", "orchid", "seagreen", "sienna", "darkkhaki")
    output_file(os.path.join('./res/plots', name.lower() + '.html'), title=name)
    s = figure(width=720, height=int(360 * (y_range[1] - y_range[0]) / 900), title="Performance",
               x_axis_label='episodes', y_axis_label='score', y_range=y_range)
    for id in sorted(summaries.keys(), key=lambda x: ids.index(x)):
        ii = ids.index(id)
        s.line(episodes[:len(summaries[id][1])], summaries[id][1], line_color=COLORS[ii], line_width=2, line_alpha=0.75,
               legend_label=legend[ids.index(id)])
        s.varea(episodes[:len(summaries[id][1])], summaries[id][0], summaries[id][2], fill_color=COLORS[ii],
                fill_alpha=0.25)
    # s.legend.location = "top_left"
    s.legend.location = "bottom_right"
    save(gridplot([[s]]))

settings = {}
folders = []
for fname in ('gradient', 'imitation', 'planet'):
    for folder in listdir(f"../{fname}/res/results"):
        folders.append((folder, fname))

for results_folder, algorithm in folders:
    with open(os.path.join(results_folder, 'hyperparameters.txt'), 'r') as f:
        s = ""
        for line in f.readlines():
            if 'seed' not in line:
                if line.split(' ')[0] == 'id':
                    id = f"id := {algorithm}_{line.split(' ')[-1]}"
                    print(id.replace('\n', ''))
                    s += id
                else:
                    s += line
    try:
        with open(os.path.join(results_folder, 'metrics.json'), 'r') as f:
            metrics = json.load(f)
        if len(metrics['steps']) >= 1000 or 'data_aggregation' in id:
            hash_append(settings, s, results_folder)
    except FileNotFoundError:
        pass

combined(['planet_lov', 'planet_vanilla'], 'planet_validation',
         legend=['With Latent Overshooting', 'Without Latent Overshooting'])
combined(['planet_lov', 'gradient_ar4'], 'planet_gradient',
         legend=['CEM Planner', 'Gradient-Based Optimization'])
combined(['planet_lov', 'imitation_data_aggregation', 'imitation_policy_aggregation'], 'planet_imitation',
         legend=['CEM Planner', 'Data Aggregation', 'Policy Aggregation'], y_range=(-200, 900))
combined(['planet_lov', 'planet_latency', 'planet_latency2', 'planet_latency4'], 'latency_planet',
         legend=['no latency', '1 timestep', '2 timesteps', '4 timesteps'])
combined(['gradient_ar4', 'gradient_ar4_lat1', 'gradient_ar4_lat2', 'gradient_ar4_lat4', 'gradient_ar4_lat8'], 'latency_gradient',
         legend=['no latency', '1 timestep', '2 timesteps', '4 timesteps', '8 timesteps'])

combined(['planet_lov', 'planet_ar4', 'planet_ar8', 'planet_ar12'], 'planet_cf',
         legend=['2 timesteps', '4 timesteps', '8 timesteps', '12 timesteps'])
combined(['gradient_ar2_4andreas', 'gradient_ar4', 'gradient_ar8', 'gradient_ar12'], 'gradient_cf',
         legend=['2 timesteps', '4 timesteps', '8 timesteps', '12 timesteps'])


import httpx
import asyncio
import logging

import discord
from discord.ext import tasks
from redbot.core import Config, commands

IDENTIFIER = 4175987634255572345  # Random to this cog

ishtakar_world_id = "3f1cd819f97e"
default_server = "Ishtakar"
realm_data_url = "https://nwdb.info/server-status/data.json"

default_guild = {
    "default_realm": "Ishtakar",
    "server_channel": None,
}

logger = logging.getLogger("red.psykzz.cogs")
logger.setLevel(logging.DEBUG)

class ServerStatus(commands.Cog):
    "Provider server status"

    def __init__(self, bot):
        self.bot = bot

        self.config = Config.get_conf(
            self, identifier=IDENTIFIER, force_registration=True
        )
        self.config.register_guild(**default_guild)

        self.refresh_queue_data.start()

    def cog_unload(self):
        self.refresh_queue_data.cancel()

    @tasks.loop(minutes=5.0)
    async def refresh_queue_data(self):
        logger.info("Starting queue task")
        try:
            self.queue_data = await self.get_queue_data(worldId=None)
            await self.update_monitor_channels()
        except Exception:
            logger.exception("Error in task")
        logger.info("Finished queue task")

    async def get_queue_data(self, worldId=ishtakar_world_id):
        """Refresh data from remote data"""
        try:
            extra_qs = f"worldId={worldId}" if worldId else ""
            response = await http_get(
                f"https://nwdb.info/server-status/servers.json?{extra_qs}"
            )
            if not response.get("success"):
                logger.error("Failed to get server status data")
                return
            servers = response.get("data", {}).get("servers", [])
            return {
                self.parse_server(server).get("worldName"): self.parse_server(server)
                for server in servers
            }
        except Exception:
            logger.exception("Exception while downloading new data")

    def parse_server(self, server):
        (
            connectionCountMax,
            connectionCount,
            queueCount,
            queueTime,
            worldName,
            worldSetName,
            region,
            status,
            active,
            worldId,
            a,
            b,
        ) = server
        return {
            "connectionCountMax": connectionCountMax,
            "connectionCount": connectionCount,
            "queueCount": queueCount,
            "queueTime": queueTime,
            "worldName": worldName,
            "worldSetName": worldSetName,
            "region": region,
            "status": status,
            "active": active,
            "worldId": worldId,
            "a-val": a,
            "b-val": b,
        }

    async def get_guild_monitor_channel(self, guild):
        guild_config = self.config.guild(guild)
        channel_id = await guild_config.server_channel()
        realm_name = await guild_config.default_realm()

        # Check if the channel is valid
        if not channel_id or channel_id == "0":
            logging.warn(f"Skipping {guild}...")
            return

        # If the channel doesn't exist, reset configuration and return
        channel = self.bot.get_channel(channel_id)
        if not channel:
            await guild_config.server_channel.set(None)
            return
        return channel

    async def update_guild_channel(self, guild):
        logger.info(f"Updating guild {guild}...")
        channel = await self.get_guild_monitor_channel(guild)

        server_status = await self.get_server_status(realm_name)
        if not server_status:
            return

        new_channel_name = server_status.split("-")[1]

        # Avoid updates if the name matches
        if channel.name == new_channel_name:
            return
        await channel.edit(name=new_channel_name)

    async def update_monitor_channels(self):
        # iterate through bot discords and get the guild config
        for guild in self.bot.guilds:
            self.update_guild_channel(guild)

    async def get_server_status(self, server_name, data=None):
        if not data:
            data = self.queue_data
        server_data = data.get(server_name)
        if not server_data:
            return "No server data available - Loading data..."

        online = server_data.get("connectionCount", -1)
        max_online = server_data.get("connectionCountMax", -1)
        in_queue_raw = int(server_data.get("queueCount", -1))
        in_queue = in_queue_raw if in_queue_raw > 1 else 0
        status = server_data.get("status", -1)
        if status == 4:
            return f"{server_name}: {online}/{max_online} Offline - Server maintenance"
        return f"{server_name}: {online}/{max_online} Online - {in_queue} in queue."

    async def get_world_id(self, server_name):
        if not self.queue_data:
            return
        server_data = self.queue_data.get(server_name)
        if not server_data:
            return
        return server_data.get("worldId")

    @commands.command()
    async def queue(self, ctx, server: str = None):
        "Get current queue information"

        if ctx.guild and server is None:
            guild_config = self.config.guild(ctx.guild)
            server = await guild_config.default_realm()

        if not server:
            await ctx.send("You must provide a server in DMs. `.queue <server>`")
            return

        worldId = await self.get_world_id(server)
        data = await self.get_queue_data(worldId=worldId)
        msg = await self.get_server_status(server, data)
        await ctx.send(msg)

    @commands.command()
    @commands.guild_only()
    @commands.admin_or_permissions(manage_channels=True)
    async def monitor(self, ctx, voice_channel: discord.VoiceChannel = None):
        "Start updating a channel wth the current realm status"

        # Check if the bot has permission to the channel
        bot_perms = voice_channel.permissions_for(ctx.me)
        if not bot_perms.manage_channels:
            await ctx.send(f'I require the "Manage Channels" permission for {voice_channel.mention} to execute that command.')
            return

        guild_config = self.config.guild(ctx.guild)
        await guild_config.server_channel.set(voice_channel.id if voice_channel else None)
        if voice_channel:
            await ctx.send(f"Setup {voice_channel} as the monitor channel.")
        else:
            await ctx.send(f"Disabled monitor channel.")

    @commands.command()
    @commands.guild_only()
    @commands.bot_has_permissions(manage_channels=True)
    @commands.admin_or_permissions(manage_channels=True)
    async def forcemonitor(self, ctx):
        "Force an update of the monitor voice channel wth the current realm status"

        voice_channel = await self.get_guild_monitor_channel(ctx.guild)
        bot_perms = voice_channel.permissions_for(ctx.me)
        if not bot_perms.manage_channels:
            await ctx.send(f'I require the "Manage Channels" permission for {voice_channel.mention} to execute that command.')
            return

        await self.update_guild_channel(ctx.guild)
        await ctx.send("Forced monitor channel update.")

    @commands.command()
    @commands.guild_only()
    @commands.admin_or_permissions(manage_channels=True)
    async def queueset(self, ctx, server: str = None):
        "Set the default server for this discord server"
        guild_config = self.config.guild(ctx.guild)

        if server is None:
            realm = await guild_config.default_realm()
            await ctx.send(f"Current server: '{realm}'.")
            return

        server_data = self.queue_data.get(server)
        if not server_data:
            await ctx.send(f"Can't find '{server}' in the server list.")
            return

        await guild_config.default_realm.set(server)
        await ctx.send(f"Server updated to '{server}'.")

async def http_get(url):
    max_attempts = 3
    attempt = 0
    while (
        max_attempts > attempt
    ):  # httpx doesn't support retries, so we'll build our own basic loop for that
        try:
            async with httpx.AsyncClient() as client:
                r = await client.get(url, headers={"user-agent": "psykzz-cogs/1.0.0"})
            if r.status_code == 200:
                return r.json()
            else:
                attempt += 1
            await asyncio.sleep(5)
        except (httpx._exceptions.ConnectTimeout, httpx._exceptions.HTTPError):
            attempt += 1
            await asyncio.sleep(5)
            pass


import numpy as np
import os
try:
    import welib.fastlib.fastlib as fastlib
except:
    import fastlib

def CPLambdaExample():
    """ Example to determine the CP-CT Lambda Pitch matrices of a turbine.
    This scrip uses the function CPCT_LambdaPitch which basically does the same as ParametricExample
    above.
    """
    ref_dir          = 'NREL5MW/'   # Folder where the fast input files are located (will be copied)
    main_file        = 'Main_Onshore_OF2.fst'  # Main file in ref_dir, used as a template
    FAST_EXE         = 'NREL5MW/OpenFAST2_x64s_ebra.exe' # Location of a FAST exe (and dll)

    # --- Computing CP and CT matrices for range of lambda and pitches
    Lambda = np.linspace(0.1,10,3)
    Pitch  = np.linspace(-10,10,4)

    CP,CT,Lambda,Pitch,MaxVal,result = fastlib.CPCT_LambdaPitch(ref_dir,main_file,Lambda,Pitch,fastExe=FAST_EXE,ShowOutputs=False,nCores=4,TMax=10)

    print('CP max',MaxVal)

    # --- Plotting matrix of CP values
    from mpl_toolkits.mplot3d import Axes3D
    from matplotlib import cm
    import matplotlib.pyplot as plt
    fig = plt.figure()
    ax = fig.gca(projection='3d')
    LAMBDA, PITCH = np.meshgrid(Lambda, Pitch)
    CP[CP<0]=0
    surf = ax.plot_surface(LAMBDA, PITCH, np.transpose(CP), cmap=cm.coolwarm, linewidth=0, antialiased=True,alpha=0.8)
    ax.scatter(MaxVal['lambda_opt'],MaxVal['pitch_opt'],MaxVal['CP_max'],c='k',marker='o',s=20)
    fig.colorbar(surf, shrink=0.5, aspect=5)
    plt.show()

if __name__=='__main__':
    CPLambdaExample()


from typing import Final

import sqlalchemy as sa
from sqlalchemy.dialects.postgresql import JSONB

from ._common import column_created_datetime, column_modified_datetime
from .base import metadata

# Intentionally includes the term "SECRET" to avoid leaking this value on a public domain
VENDOR_SECRET_PREFIX: Final[str] = "OSPARC_VARIABLE_VENDOR_SECRET_"

services_vendor_secrets = sa.Table(
    "services_vendor_secrets",
    #
    # - A secret is an environment value passed to the service at runtime
    # - A vendor can associate secrets (e.g. a license code) to any of the services it owns
    # - secrets_map
    #   - keys should be prefixed with OSPARC_VARIABLE_VENDOR_SECRET_ (can still normalize on read)
    #   - values might be encrypted
    #
    metadata,
    sa.Column(
        "service_key",
        sa.String,
        doc="A single environment is allowed per service",
    ),
    sa.Column(
        "service_base_version",
        sa.String,
        doc="Defines the minimum version (included) from which these secrets apply",
    ),
    sa.Column(
        "product_name",
        sa.String,
        sa.ForeignKey(
            "products.name",
            name="fk_services_name_products",
            onupdate="CASCADE",
            ondelete="CASCADE",
        ),
        # NOTE: since this is part of the primary key this is required
        # NOTE: an alternative would be to not use this as a primary key
        server_default="osparc",
        doc="Product Identifier",
    ),
    sa.Column(
        "secrets_map",
        JSONB,
        nullable=False,
        server_default=sa.text("'{}'::jsonb"),
        doc="Maps OSPARC_VARIABLE_VENDOR_SECRET_* identifiers to a secret value (could be encrypted) "
        "that can be replaced at runtime if found in the compose-specs",
    ),
    # TIME STAMPS ----
    column_created_datetime(timezone=True),
    column_modified_datetime(timezone=True),
    # CONSTRAINTS --
    sa.ForeignKeyConstraint(
        ["service_key", "service_base_version"],
        ["services_meta_data.key", "services_meta_data.version"],
        onupdate="CASCADE",
        ondelete="CASCADE",
        # NOTE: this might be a problem: if a version in the metadata is deleted,
        # all versions above will take the secret_map for the previous one.
    ),
    sa.PrimaryKeyConstraint(
        "service_key",
        "service_base_version",
        "product_name",
        name="services_vendor_secrets_pk",
    ),
)


import pyomo.environ as pe
from lci import LifeCycleInventory
from superstructure import Superstructure
from utils.properties import molar_weight
from utils.utils import sum_rule
from utils.save_results import ResultManager
from utils.solve_model import Solver
import time as time
import pickle

def range_len(start, stop, len):
    step = (stop - start) / (len - 1)
    list = []
    i = start
    while i < stop:
        list.append(i)
        i += step
    list.append(stop)
    return list

def rep_solve_tcm(x):
    lci = LifeCycleInventory('millgas2what')

    lci.model = pe.ConcreteModel('millgas2what')

    scale = 1000000000

    'Ströme des Stahlwerks als Parameter'

    lci.model.cog_steelMill = pe.Param(initialize=39700000000 / scale)  # in kg, scaled
    lci.model.bfg_steelMill = pe.Param(initialize=1740550000000 / scale)  # in kg, scaled
    lci.model.electricity_steelMill = pe.Param(initialize=-2298568545000 / scale)
    lci.model.heat_steelMill = pe.Param(initialize=-894991475000 / scale)

    'Definition der Verbindungspunkte'

    connect_list = ['Mill gas COG [kg]', 'Mill gas BFG/BOFG [kg]', 'Electricity [MJ]', 'Heat [MJ]']

    connector_lp = {}
    lci.model.connect_lp = pe.Var(connect_list, initialize=0, bounds=(-10000, 10000))

    for c in connect_list:
        connector_lp[c] = lci.model.connect_lp[c]

    'Gesamtbilanzen'

    lci.model.cog_balance = pe.Constraint(
        expr=0 == - lci.model.cog_steelMill + lci.model.connect_lp['Mill gas COG [kg]'])
    lci.model.bfg_balance = pe.Constraint(
        expr=0 == - lci.model.bfg_steelMill + lci.model.connect_lp['Mill gas BFG/BOFG [kg]'])
    lci.model.electricity_balance = pe.Constraint(
        expr=0 == - lci.model.electricity_steelMill + lci.model.connect_lp['Electricity [MJ]'])
    lci.model.heat_balance = pe.Constraint(expr=0 == - lci.model.heat_steelMill + lci.model.connect_lp['Heat [MJ]'])

    'Ab hier wird das Modell mit der LCI-Klasse zusammengebaut'

    lci.import_from_excel('Life Cycle Inventory_v19.xlsx', 'A-Matrix', 'End of life')
    lci.set_up_lp(scale)
    lci.import_connector(connector_lp)  # Durch deaktivieren dieser Zeile wird nur die Chem. Ind. betrachtet

    # lci.activate_scenario('Electricity Today')
    # lci.activate_scenario('Electricity Best Case')
    lci.activate_scenario('Electricity user-defined', x)

    lci.activate_scenario('Separation GDP')  # Schaltet alle linearen Prozesse für Hüttengastrennung aus

    # lci.deactivate_process('CARBON DIOXIDE as emission to air')

    # lci.deactivate_process('AMMONIA FROM NATURAL GAS BY STEAM REFORMING BY ICI "AMV" PROCESS incl CO2 capture')

    lci.activate_scenario('CCU high TRL only')
    # lci.activate_scenario('No high TRL CCU')
    lci.deactivate_process('CARBON DIOXIDE - ammonia plant')
    lci.deactivate_process('Water gas shift reaction')
    lci.lp.ammonia_constraint = pe.Constraint(
        expr=lci.lp.s['AMMONIA FROM NATURAL GAS BY STEAM REFORMING BY ICI "AMV" PROCESS incl CO2 capture'] - 228.9 <= 0)
    lci.deactivate_process('TDI - neue Route_v2 exklusive Methylformate production')
    lci.deactivate_process('Polycarbonate - neue Route')
    lci.deactivate_process('Methylformate productionaus TDI neue Route v2')

    lci.construct_demand_constraints()
    lci.construct_objective()

    lci.model.add_component('lp', lci.lp)

    'Lösen und Ergebnisse darstellen'

    # lci.model.pprint()

    solver = Solver()
    solver.solve_lp(lci, 'glpk')
    # solver.test_feasibility() # Muss an neue Demand-Constraints angepasst werden

    results = {'x': x, 'z': pe.value(lci.objective) * lci.scale}

    return results

def rep_solve_gdp(x):
    """ Separation System / Flowsheet construction """

    # COG Separation

    s = Superstructure("Combined Model")

    s.initial_stream(31, 0, 300, 1, 'COG')
    s.create_unit('Compressor', 'C3', 31, 32)
    s.create_unit('PSA', 'PSA3', 32, 33, 34, 'H2')

    s.create_unit('Splitter', 'S5', 34, 35, 36)

    s.create_streams(37)
    s.model.p_35 = pe.Constraint(expr=s.model.p[37] == 1)
    s.model.t_35 = pe.Constraint(expr=s.model.t[37] == 300)
    s.model.cc_35 = pe.Constraint(expr=1 == sum_rule(s.streams[37].y, s.streams[37].substances))

    s.mix_streams('M3', 36, 37, 38)

    s.create_unit('Heat Exchanger', 'HE5', 38, 39)

    s.create_disjunct_reactor('methane_reforming', 'por', 'R2_POR', 39, 40, 'POR')
    s.create_disjunct_reactor('methane_reforming', 'cdr', 'R2_CDR', 39, 40, 'CDR')

    s.model.cdr.y_37 = pe.Constraint(expr=s.model.y[37, 'CO2'] == 1)
    s.model.cdr.add = pe.Constraint(expr=s.model.n[37] == 1 * s.model.n[36] * s.model.y[36, 'CH4'])
    s.model.cdr.q_constr = pe.Constraint(expr=s.model.q['R2_POR'] == 0)

    s.model.por.y_37 = pe.Constraint(expr=s.model.y[37, 'O2'] == 1)
    s.model.por.add = pe.Constraint(expr=s.model.n[37] == 0.48 * s.model.n[36] * s.model.y[36, 'CH4'])
    s.model.por.q_constr = pe.Constraint(expr=s.model.q['R2_CDR'] == 0)

    # B(O)FG Separation

    s.initial_stream(1, 0, 300, 1, 'B(O)FG')

    s.create_unit('Heat Exchanger', 'HE1', 1, 2)

    s.create_unit('TSA', 'TSA1', 2, 3, 4, 'CO')

    s.create_unit('Splitter', 'S1', 3, 20, 21)

    s.initial_stream(22, 0, 400, 1, {'H2O': 1})
    s.model.add_h2o = pe.Constraint(expr=s.model.n[22] == s.model.n[21])

    s.mix_streams('M1', 21, 22, 23)

    s.create_unit('Heat Exchanger', 'HE3', 23, 24)

    s.create_reactor('R1', 24, 25, 'WGSR')

    s.mix_streams('M2', 25, 4, 5)

    s.create_unit('Heat Exchanger', 'HE4', 5, 6)

    s.create_unit('Compressor', 'C1', 6, 7)

    s.create_disjunct_unit('co2_bofg', 'cca_1', 'CCA', 'CCA1', 7, 8, 9, 'CO2')
    s.create_disjunct_unit('co2_bofg', 'psa_1', 'PSA', 'PSA1', 7, 8, 9, 'CO2')
    s.model.psa_1.q_cca = pe.Constraint(expr=s.model.q['CCA1'] == 0)
    s.model.psa_1.t6_constraint = pe.Constraint(expr=s.model.t[6] >= 273.15 + 30)

    # s.model.t6_constraint = pe.Constraint(expr=s.model.t[6] >= 273.15 + 30)
    # s.create_unit('PSA', 'PSA1', 7, 8, 9, 'CO2')

    # s.create_unit('CCA', 'CCA1', 7, 8, 9, 'CO2')

    s.create_unit('Compressor', 'C2', 9, 10)

    s.create_unit('PSA', 'PSA2', 10, 11, 12, 'H2')

    scale = 1000000000

    'Ströme des Stahlwerks als Parameter'

    s.model.cog_steelMill = pe.Param(initialize=39700000000 / scale)  # in kg, scaled
    s.model.bfg_steelMill = pe.Param(initialize=1740550000000 / scale)  # in kg, scaled

    s.model.electricity_steelMill = pe.Param(initialize=-2298568545000 / scale)
    s.model.heat_steelMill = pe.Param(initialize=-894991475000 / scale)

    s.connect_list = ['Mill gas COG [kg]', 'Hydrogen (H2) [kg]', 'Electricity [MJ]', 'Heat [MJ]', 'SYNTHESIS GAS (1:1)',
                      'SYNTHESIS GAS (2:1)', 'Carbon dioxide (CO2) [kg]', 'Methane (CH4) [kg]', 'Oxygen (O2) [kg]',
                      'Mill gas BFG/BOFG [kg]', 'Carbon monoxide (CO) [kg]', 'CO2 to atm [kg]', 'STEAM [kg]'
                      ]

    connector_lp = {}
    s.model.connect_lp = pe.Var(s.connect_list, initialize=0, bounds=(-10000, 10000))

    for c in s.connect_list:
        connector_lp[c] = s.model.connect_lp[c]

    # Hier kann eingestellt werden, welcher Anteil der Hüttengase verwertet werden darf
    # x = 0.0001  # ACHTUNG! x = 0 führt zu infeasible
    # s.model.force_cog = pe.Constraint(expr=s.model.n[31] <= x * s.model.cog_steelMill / molar_weight('COG'))
    # s.model.force_bfg = pe.Constraint(expr=s.model.n[1] <= x * s.model.bfg_steelMill / molar_weight('B(O)FG'))

    s.model.cog_balance = pe.Constraint(
        expr=0 == - s.model.cog_steelMill + s.model.connect_lp['Mill gas COG [kg]'] + s.model.n[31] * molar_weight(
            'COG'))
    s.model.bfg_balance = pe.Constraint(
        expr=0 == - s.model.bfg_steelMill + s.model.connect_lp['Mill gas BFG/BOFG [kg]'] + s.model.n[1] * molar_weight(
            'B(O)FG'))
    s.model.el_balance = pe.Constraint(expr=0 == - s.model.connect_lp['Electricity [MJ]'] - sum_rule(s.model.w,
                                                                                s.model.workSet) + s.model.electricity_steelMill)
    s.model.heat_balance = pe.Constraint(
        expr=0 == - s.model.connect_lp['Heat [MJ]'] - sum_rule(s.model.q, s.model.heatSet) + s.model.heat_steelMill)
    s.model.co2_atm_balance = pe.Constraint(
        expr=0 == s.model.connect_lp['CO2 to atm [kg]'] - s.model.n[12] * molar_weight('CO2') * (
                s.model.y[12, 'CO2'] + s.model.y[12, 'CO']))
    # s.model.co2_atm_balance = pe.Constraint(expr=0 == s.model.connect_lp['CO2 to atm [kg]'] - s.model.n[12] * molar_weight('CO2') * (s.model.y[12, 'CO2'] + s.model.y[12, 'CO']) * 0.05)
    # s.model.co2_atm_balance = pe.Constraint(expr=0 == s.model.connect_lp['CO2 to atm [kg]'])
    s.model.steam_balance = pe.Constraint(
        expr=0 == s.model.connect_lp['STEAM [kg]'] + s.model.n[22] * molar_weight('H2O'))
    s.model.co_balance = pe.Constraint(
        expr=0 == s.model.connect_lp['Carbon monoxide (CO) [kg]'] - s.model.n[20] * molar_weight('CO'))
    s.model.ch4_balance = pe.Constraint(
        expr=s.model.connect_lp['Methane (CH4) [kg]'] == s.model.n[35] * s.model.y[35, 'CH4'] * molar_weight('CH4'))
    s.model.h2_balance = pe.Constraint(
        expr=s.model.connect_lp['Hydrogen (H2) [kg]'] == s.model.n[33] * molar_weight('H2') + s.model.n[
            11] * molar_weight('H2'))

    # s.model.cdr.syngas11_balance = pe.Constraint(expr=s.model.connect_lp['SYNTHESIS GAS (1:1)'] == s.model.n[40] * molar_weight({'H2': 0.5, 'CO': 0.5}))  # N2 wird als SynGas angenommen
    s.model.cdr.syngas11_balance = pe.Constraint(expr=s.model.connect_lp['SYNTHESIS GAS (1:1)'] == s.model.n[40] * (
                s.model.y[40, 'H2'] * molar_weight('H2') + s.model.y[40, 'CO'] * molar_weight('CO')))
    s.model.por.syngas11_balance = pe.Constraint(expr=s.model.connect_lp['SYNTHESIS GAS (1:1)'] == 0)

    # s.model.por.syngas21_balance = pe.Constraint(expr=s.model.connect_lp['SYNTHESIS GAS (2:1)'] == s.model.n[40] * molar_weight({'H2': 0.67, 'CO': 0.33}))  # N2 wird als SynGas angenommen
    s.model.por.syngas21_balance = pe.Constraint(expr=s.model.connect_lp['SYNTHESIS GAS (2:1)'] == s.model.n[40] * (
                s.model.y[40, 'H2'] * molar_weight('H2') + s.model.y[40, 'CO'] * molar_weight('CO')))
    s.model.cdr.syngas21_balance = pe.Constraint(expr=s.model.connect_lp['SYNTHESIS GAS (2:1)'] == 0)

    s.model.cdr.co2_balance = pe.Constraint(
        expr=s.model.connect_lp['Carbon dioxide (CO2) [kg]'] == - s.model.n[37] * molar_weight('CO2') + s.model.n[
            8] * molar_weight('CO2'))
    s.model.por.co2_balance = pe.Constraint(
        expr=s.model.connect_lp['Carbon dioxide (CO2) [kg]'] == s.model.n[8] * molar_weight('CO2'))
    # s.model.cdr.co2_balance = pe.Constraint(expr=s.model.connect_lp['Carbon dioxide (CO2) [kg]'] == 0)
    # s.model.por.co2_balance = pe.Constraint(expr=s.model.connect_lp['Carbon dioxide (CO2) [kg]'] == 0)
    s.model.por.o2_balance = pe.Constraint(
        expr=s.model.connect_lp['Oxygen (O2) [kg]'] == - s.model.n[37] * molar_weight('O2'))
    s.model.cdr.o2_balance = pe.Constraint(expr=s.model.connect_lp['Oxygen (O2) [kg]'] == 0)

    # s.model.no_co2_sep = pe.Constraint(expr=s.model.n[8] * molar_weight('CO2') == 147)
    # s.model.no_h2_sep = pe.Constraint(expr=s.model.n[11] == 0)

    """ Set up TCM """

    lci = LifeCycleInventory('millgas2what')

    lci.import_from_excel('Life Cycle Inventory_v19.xlsx', 'A-Matrix', 'End of life')
    lci.set_up_lp(scale)
    lci.import_connector(connector_lp)  # Durch deaktivieren dieser Zeile wird nur die Chem. Ind. betrachtet

    # lci.activate_scenario('Electricity Today')
    # lci.activate_scenario('Electricity Best Case')
    lci.activate_scenario('Electricity user-defined', x)

    lci.activate_scenario('Separation GDP')  # Schaltet alle linearen Prozesse für Hüttengastrennung aus

    # lci.deactivate_process('CARBON DIOXIDE as emission to air')

    # lci.deactivate_process('AMMONIA FROM NATURAL GAS BY STEAM REFORMING BY ICI "AMV" PROCESS incl CO2 capture')

    lci.activate_scenario('CCU high TRL only')
    lci.deactivate_process('CARBON DIOXIDE - ammonia plant')
    lci.deactivate_process('Water gas shift reaction')
    lci.lp.ammonia_constraint = pe.Constraint(
        expr=lci.lp.s['AMMONIA FROM NATURAL GAS BY STEAM REFORMING BY ICI "AMV" PROCESS incl CO2 capture'] - 228.9 <= 0)
    lci.deactivate_process('TDI - neue Route_v2 exklusive Methylformate production')
    lci.deactivate_process('Polycarbonate - neue Route')
    lci.deactivate_process('Methylformate productionaus TDI neue Route v2')

    lci.construct_demand_constraints()
    lci.construct_objective()

    s.import_lci(lci)

    s.create_disjunctions()

    """ Solve overall model """

    solver = Solver()

    """ Save values in dict"""

    ind_var = {}
    c_val = {}
    rec_val = {}
    heat_val = {}
    el_val = {}
    s_dict = {}
    cog_dict = {}
    bfg_dict = {}

    # LCI-Prozesse, die verfolgt werden sollen
    s_dict_list = ['CARBON DIOXIDE as emission to air', 'INC Carbon monoxide',
                   'treatment of blast furnace gas, in power plant, DE', 'Verbrennung COG in BHKW',
                   'AMMONIA FROM NATURAL GAS BY STEAM REFORMING BY ICI "AMV" PROCESS incl CO2 capture',
                   'CARBON DIOXIDE - air capture', 'Electricity, user-defined']

    try:
        solver.solve_gdp(s)
        obj = pe.value(s.objective)
        for d in s.disjuncts.keys():
            ind_var[d] = pe.value(s.disjuncts[d].indicator_var)
        for c in s.connect_list:
            c_val[c] = pe.value(connector_lp[c]) * lci.scale
        for c in s.model.zetaSet:
            rec_val[c] = pe.value(s.model.zeta[c])
        for c in s.model.heatSet:
            heat_val[c] = pe.value(s.model.q[c]) * lci.scale
        for c in s.model.workSet:
            el_val[c] = pe.value(s.model.w[c]) * lci.scale
        for i in s_dict_list:
            s_dict[i] = pe.value(s.model.utilization.s[i]) * lci.scale
        'Diagramme für COG / B(O)FG'
        cog_dict['H2'] = pe.value(s.model.n[33]) * molar_weight({'H2': 1}) * lci.scale
        cog_dict['N2'] = pe.value(s.model.n[35]) * pe.value(s.model.y[35, 'N2']) * molar_weight({'N2': 1}) * lci.scale

        bfg_dict['H2'] = pe.value(s.model.n[11]) * molar_weight({'H2': 1}) * lci.scale
        cog_dict['N2'] = pe.value(s.model.n[12]) * pe.value(s.model.y[12, 'N2']) * molar_weight({'N2': 1}) * lci.scale
        bfg_dict['CO2'] = pe.value(s.model.n[8]) * molar_weight({'CO2': 1}) * lci.scale

    except ValueError:
        obj = 0

    return {'x': x, 'z': obj * lci.scale, 'i': ind_var, 'c': c_val, 'rec': rec_val, 'q': heat_val, 'w': el_val,
            's': s_dict, 'cog': cog_dict, 'bfg': bfg_dict}

x_vector = range_len(0.002, 0.2, 30)

results_tcm = {}
results_gdp = {}

n = 0
while n < len(x_vector):
    results_tcm[n] = rep_solve_tcm(x_vector[n])
    print('TCM solved', n + 1)
    results_gdp[n] = rep_solve_gdp(x_vector[n])
    print('GDP solved', n + 1)
    n += 1

results = {'tcm': results_tcm, 'gdp': results_gdp}

def save_object(obj, filename):
    with open(filename, 'wb') as output:  # Overwrites any existing file.
        pickle.dump(obj, output, pickle.HIGHEST_PROTOCOL)

# name = input("enter file name")
save_object(results, '20200825_v19_tcm')


'''
  1. import opencv
  2. load image
  3. load model
  4. adjuct image gray
  5. Check and mark face
  3. create window
  4. show image
  5. pause window
  6. close window
'''

import numpy as np
import cv2
# print(cv2.__version__)

# Load image
img = cv2.imread("./images/ufc.jpeg")

# Load model
face_cascade = cv2.CascadeClassifier(
    './data/haarcascade_frontalface_default.xml')

# Adjust image gray
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

# find faces
faces = face_cascade.detectMultiScale(
    gray,
    # scaleFactor=1.15,
    # minNeighbors=5,
    # minSize=(5, 5),
    #   flags = cv2.HAAR_SCALE_IMAGE
)

# Mark faces
for(x, y, w, h) in faces:
    # image, location, size, color, line width
    cv2.rectangle(img, (x, y), (x+w, y+w), (0, 255, 0), 2)

# Create widnow
cv2.namedWindow("faceImage")

# Show image
cv2.imshow('ufc show', img)

#
cv2.waitKey(0)

cv2.destroyAllWindow()


import os

os.chdir("/home/ghiggi/Projects/deepsphere-weather")
import sys

sys.path.append("../")
import shutil
import argparse
import dask
import glob
import time
import torch
import zarr
import numpy as np
import xarray as xr

## DeepSphere-Weather
from modules.utils_config import read_config_file
from modules.utils_config import get_model_settings
from modules.utils_config import get_training_settings
from modules.utils_config import get_ar_settings
from modules.utils_config import get_dataloader_settings
from modules.utils_config import check_same_dict
from modules.utils_config import get_pytorch_model
from modules.utils_config import set_pytorch_settings
from modules.utils_config import load_pretrained_model
from modules.utils_config import print_tensor_info
from modules.utils_io import get_ar_model_tensor_info
from modules.predictions_autoregressive import AutoregressivePredictions

## Functions within AutoregressivePredictions
from modules.dataloader_autoregressive import remove_unused_Y
from modules.dataloader_autoregressive import get_aligned_ar_batch
from modules.dataloader_autoregressive import AutoregressiveDataset
from modules.dataloader_autoregressive import AutoregressiveDataLoader
from modules.utils_autoregressive import check_ar_settings
from modules.utils_autoregressive import check_input_k
from modules.utils_autoregressive import check_output_k
from modules.utils_io import _get_feature_order
from modules.utils_zarr import check_chunks
from modules.utils_zarr import check_rounding
from modules.utils_zarr import rechunk_Dataset
from modules.utils_zarr import write_zarr
from modules.utils_torch import check_device
from modules.utils_torch import check_pin_memory
from modules.utils_torch import check_asyncronous_gpu_transfer
from modules.utils_torch import check_prefetch_in_gpu
from modules.utils_torch import check_prefetch_factor
from modules.utils_swag import bn_update

## Project specific functions
import modules.my_models_graph as my_architectures

## Side-project utils (maybe migrating to separate packages in future)
from modules.xscaler import LoadScaler
from modules.xscaler import SequentialScaler

# -------------------------------------------------------------------------.

data_dir = "/ltenas3/DeepSphere/data/preprocessed_ds/ERA5_HRES"
model_dir = "/data/weather_prediction/experiments_GG/new/RNN-AR6-UNetSpherical-Healpix_400km-Graph_knn-k20-MaxAreaPooli/"

# -------------------------------------------------------------------------.
# Read config file
cfg_path = os.path.join(model_dir, "config.json")
cfg = read_config_file(fpath=cfg_path)
# Some special options to adjust for prediction
cfg["dataloader_settings"]["autotune_num_workers"] = False
cfg["training_settings"]["gpu_training"] = True  # to run prediction in GPU if possible
##------------------------------------------------------------------------.
### Retrieve experiment-specific configuration settings
model_settings = get_model_settings(cfg)
ar_settings = get_ar_settings(cfg)
training_settings = get_training_settings(cfg)
dataloader_settings = get_dataloader_settings(cfg)
dataloader_settings["num_workers"] = 10

##------------------------------------------------------------------------.
#### Load Zarr Datasets
data_sampling_dir = os.path.join(data_dir, cfg["model_settings"]["sampling_name"])

data_dynamic = xr.open_zarr(
    os.path.join(data_sampling_dir, "Data", "dynamic", "time_chunked", "dynamic.zarr")
)
data_bc = xr.open_zarr(
    os.path.join(data_sampling_dir, "Data", "bc", "time_chunked", "bc.zarr")
)
data_static = xr.open_zarr(os.path.join(data_sampling_dir, "Data", "static.zarr"))

# - Select dynamic features
# data_dynamic = data_dynamic[['z500','t850']]

##------------------------------------------------------------------------.
### Prepare static data
# - Keep land-surface mask as it is

# - Keep sin of latitude and remove longitude information
data_static = data_static.drop(["sin_longitude", "cos_longitude"])

# - Scale orography between 0 and 1 (is already left 0 bounded)
data_static["orog"] = data_static["orog"] / data_static["orog"].max()

# - One Hot Encode soil type
# ds_slt_OHE = xscaler.OneHotEnconding(data_static['slt'])
# data_static = xr.merge([data_static, ds_slt_OHE])
# data_static = data_static.drop('slt')

# - Load static data
data_static = data_static.load()

##------------------------------------------------------------------------.
#### Define scaler to apply on the fly within DataLoader
# - Load scalers
dynamic_scaler = LoadScaler(
    os.path.join(data_sampling_dir, "Scalers", "GlobalStandardScaler_dynamic.nc")
)
bc_scaler = LoadScaler(
    os.path.join(data_sampling_dir, "Scalers", "GlobalStandardScaler_bc.nc")
)
# # - Create single scaler
scaler = SequentialScaler(dynamic_scaler, bc_scaler)

##------------------------------------------------------------------------.
### Define pyTorch settings (before PyTorch model definition)
# - Here inside is eventually set the seed for fixing model weights initialization
# - Here inside the training precision is set (currently only float32 works)
device = set_pytorch_settings(training_settings)

##------------------------------------------------------------------------.
## Retrieve dimension info of input-output Torch Tensors
tensor_info = get_ar_model_tensor_info(
    ar_settings=ar_settings,
    data_dynamic=data_dynamic,
    data_static=data_static,
    data_bc=data_bc,
)
print_tensor_info(tensor_info)

# Check that tensor_info match between model training and now
check_same_dict(model_settings["tensor_info"], tensor_info)

##------------------------------------------------------------------------.
### Define the model architecture
model = get_pytorch_model(module=my_architectures, model_settings=model_settings)

###-----------------------------------------------------------------------.
## Load a pre-trained model
load_pretrained_model(model=model, model_dir=model_dir)

###-----------------------------------------------------------------------.
### Transfer model to the device (i.e. GPU)
model = model.to(device)

###-----------------------------------------------------------------------.
## AutoregressivePredictions arguments
forecast_reference_times = np.datetime64("2016-12-26T23:00:00.000000000")
forecast_reference_times1 = np.datetime64("2016-06-26T23:00:00.000000000")
forecast_reference_times = [forecast_reference_times, forecast_reference_times1]
ar_iterations = 2 * 365 * 4
ar_iterations = 20
batch_size = 32
ar_blocks = None
forecast_zarr_fpath = None
num_workers = 10  # dataloader_settings['num_workers']

bc_generator = None
ar_batch_fun = get_aligned_ar_batch
scaler_transform = scaler
scaler_inverse = scaler
# Dataloader options
device = device
batch_size = batch_size  # number of forecasts per batch

prefetch_factor = dataloader_settings["prefetch_factor"]
prefetch_in_gpu = dataloader_settings["prefetch_in_gpu"]
pin_memory = dataloader_settings["pin_memory"]
asyncronous_gpu_transfer = dataloader_settings["asyncronous_gpu_transfer"]
# Autoregressive settings
input_k = ar_settings["input_k"]
output_k = ar_settings["output_k"]
forecast_cycle = ar_settings["forecast_cycle"]
stack_most_recent_prediction = ar_settings["stack_most_recent_prediction"]
# Prediction options
forecast_reference_times = forecast_reference_times
ar_blocks = ar_blocks
ar_iterations = ar_iterations  # How many time to autoregressive iterate
keep_first_prediction = True
# Save options
zarr_fpath = forecast_zarr_fpath  # None --> do not write to disk
rounding = 2  # Default None. Accept also a dictionary
compressor = "auto"  # Accept also a dictionary per variable
chunks = "auto"

# 1 Valid timestep : OK
forecast_reference_times = np.datetime64("2018-12-26T23:00:00.000000000")

### 2 (valid) timesteps --> OK
forecast_reference_times1 = np.datetime64("2018-12-26T22:00:00.000000000")
forecast_reference_times2 = np.datetime64("2018-12-26T23:00:00.000000000")
forecast_reference_times = [forecast_reference_times2, forecast_reference_times1]

## One valid, one unvalid
forecast_reference_times1 = np.datetime64("2018-12-26T23:00:00.000000000")
forecast_reference_times2 = np.datetime64("2018-12-27T00:00:00.000000000")
forecast_reference_times = [forecast_reference_times2, forecast_reference_times1]

## 1 Unvalid (future) timestep  --> OK: raise correct error
forecast_reference_times = np.datetime64("2018-12-27T00:00:00.000000000")

## 1 Unvalid timestep (past) --> OK: raise correct error
forecast_reference_times = np.datetime64("1980-01-01T07:00:00.000000000")

forecast_reference_times = np.datetime64("1970-01-01T07:00:00.000000000")

## 2 unvalid (future) timesteps  --> OK: raise correct error
forecast_reference_times1 = np.datetime64("2018-12-27T00:00:00.000000000")
forecast_reference_times2 = np.datetime64("2018-12-27T01:00:00.000000000")
forecast_reference_times = [forecast_reference_times2, forecast_reference_times1]

## 2 unvalid (past) timesteps  --> OK: raise correct error
forecast_reference_times1 = np.datetime64("1980-01-01T07:00:00.000000000")
forecast_reference_times2 = np.datetime64("1980-01-01T06:00:00.000000000")
forecast_reference_times = [forecast_reference_times2, forecast_reference_times1]

# ----
### No duplicate (unvalid) timesteps --> OK raise correct error
forecast_reference_times1 = np.datetime64("2018-12-27T00:00:00.000000000")
forecast_reference_times2 = np.datetime64("2018-12-27T00:00:00.000000000")
forecast_reference_times = [forecast_reference_times2, forecast_reference_times1]

### No duplicate (valid) timesteps --> OK raise correct error
forecast_reference_times1 = np.datetime64("2018-12-26T23:00:00.000000000")
forecast_reference_times2 = np.datetime64("2018-12-26T23:00:00.000000000")
forecast_reference_times = [forecast_reference_times2, forecast_reference_times1]

## Empty list --> OK raise correct error
forecast_reference_times = []
# ----

## AutoregressivePredictions arguments
forecast_reference_times = np.datetime64("2016-12-26T23:00:00.000000000")
forecast_reference_times1 = np.datetime64("2016-06-26T23:00:00.000000000")
forecast_reference_times = [forecast_reference_times, forecast_reference_times1]
ar_iterations = 2 * 365 * 4

dask.config.set(scheduler="synchronous")
ds_forecasts = AutoregressivePredictions(
    model=model,
    # Data
    data_dynamic=data_dynamic,
    data_static=data_static,
    data_bc=data_bc,
    scaler_transform=scaler,
    scaler_inverse=scaler,
    # Dataloader options
    device=device,
    batch_size=batch_size,  # number of forecasts per batch
    num_workers=dataloader_settings["num_workers"],
    prefetch_factor=dataloader_settings["prefetch_factor"],
    prefetch_in_gpu=dataloader_settings["prefetch_in_gpu"],
    pin_memory=dataloader_settings["pin_memory"],
    asyncronous_gpu_transfer=dataloader_settings["asyncronous_gpu_transfer"],
    # Autoregressive settings
    input_k=ar_settings["input_k"],
    output_k=ar_settings["output_k"],
    forecast_cycle=ar_settings["forecast_cycle"],
    stack_most_recent_prediction=ar_settings["stack_most_recent_prediction"],
    # Prediction options
    forecast_reference_times=forecast_reference_times,
    ar_blocks=ar_blocks,
    ar_iterations=ar_iterations,  # How many time to autoregressive iterate
    # Save options
    zarr_fpath=forecast_zarr_fpath,  # None --> do not write to disk
    rounding=2,  # Default None. Accept also a dictionary
    compressor="auto",  # Accept also a dictionary per variable
    chunks="auto",
)
print(ds_forecasts)

ds_forecasts.to_zarr("/ltenas3/DeepSphere/tmp/2ysim.zarr")

###-----------------------------------------------------------------------.
## DEBUG Code within AutoregressivePredictions
##------------------------------------------------------------------------.
## Checks arguments
device = check_device(device)
pin_memory = check_pin_memory(
    pin_memory=pin_memory, num_workers=num_workers, device=device
)
asyncronous_gpu_transfer = check_asyncronous_gpu_transfer(
    asyncronous_gpu_transfer=asyncronous_gpu_transfer, device=device
)
prefetch_in_gpu = check_prefetch_in_gpu(
    prefetch_in_gpu=prefetch_in_gpu, num_workers=num_workers, device=device
)
prefetch_factor = check_prefetch_factor(
    prefetch_factor=prefetch_factor, num_workers=num_workers
)
##------------------------------------------------------------------------.
# Check that autoregressive settings are valid
# - input_k and output_k must be numpy arrays hereafter !
input_k = check_input_k(input_k=input_k, ar_iterations=ar_iterations)
output_k = check_output_k(output_k=output_k)
check_ar_settings(
    input_k=input_k,
    output_k=output_k,
    forecast_cycle=forecast_cycle,
    ar_iterations=ar_iterations,
    stack_most_recent_prediction=stack_most_recent_prediction,
)
ar_iterations = int(ar_iterations)
##------------------------------------------------------------------------.
### Retrieve feature info of the forecast
features = _get_feature_order(data_dynamic)

##------------------------------------------------------------------------.
# Check Zarr settings
WRITE_TO_ZARR = zarr_fpath is not None
if WRITE_TO_ZARR:
    # - If zarr fpath provided, create the required folder
    if not os.path.exists(os.path.dirname(zarr_fpath)):
        os.makedirs(os.path.dirname(zarr_fpath))
    # - Set default chunks and compressors
    # ---> -1 to all optional dimensions (i..e nodes, lat, lon, ens, plevels,...)
    dims = list(data_dynamic.dims)
    dims_optional = np.array(dims)[
        np.isin(dims, ["time", "feature"], invert=True)
    ].tolist()
    default_chunks = {dim: -1 for dim in dims_optional}
    default_chunks["forecast_reference_time"] = 1
    default_chunks["leadtime"] = 1
    default_compressor = zarr.Blosc(cname="zstd", clevel=0, shuffle=2)
    # - Check rounding settings
    rounding = check_rounding(rounding=rounding, variable_names=features)
##------------------------------------------------------------------------.
# Check ar_blocks
if not isinstance(ar_blocks, (int, float, type(None))):
    raise TypeError("'ar_blocks' must be int or None.")
if isinstance(ar_blocks, float):
    ar_blocks = int(ar_blocks)
if not WRITE_TO_ZARR and isinstance(ar_blocks, int):
    raise ValueError("If 'zarr_fpath' not specified, 'ar_blocks' must be None.")
if ar_blocks is None:
    ar_blocks = ar_iterations + 1
if ar_blocks > ar_iterations + 1:
    raise ValueError("'ar_blocks' must be equal or smaller to 'ar_iterations'")
PREDICT_ar_BLOCKS = ar_blocks != (ar_iterations + 1)

##------------------------------------------------------------------------.
### Define DataLoader subset_timesteps
forecast_reference_times = check_timesteps_format(forecast_reference_times)
check_no_duplicate_timesteps(
    forecast_reference_times, var_name="forecast_reference_times"
)
forecast_reference_times.sort()  # ensure the temporal order
subset_timesteps = None
if forecast_reference_times is not None:
    if len(forecast_reference_times) == 0:
        raise ValueError(
            "If you don't want to specify specific 'forecast_reference_times', set it to None"
        )
    t_res_timedelta = np.diff(data_dynamic.time.values)[0]
    subset_timesteps = forecast_reference_times + -1 * max(input_k) * t_res_timedelta

##------------------------------------------------------------------------.
### Create training Autoregressive Dataset and DataLoader
dataset = AutoregressiveDataset(
    data_dynamic=data_dynamic,
    data_bc=data_bc,
    data_static=data_static,
    bc_generator=bc_generator,
    scaler=scaler_transform,
    # Dataset options
    subset_timesteps=subset_timesteps,
    training_mode=False,
    # Autoregressive settings
    input_k=input_k,
    output_k=output_k,
    forecast_cycle=forecast_cycle,
    ar_iterations=ar_iterations,
    stack_most_recent_prediction=stack_most_recent_prediction,
    # GPU settings
    device=device,
)

dataset[0]
self = dataset
self.subset_timesteps
self.idxs
len(self)


import pathlib

data_folder = pathlib.Path('data')

# print(data_folder.exists(), data_folder.is_dir())

def make_text(i):
    text = ""
    text += str(i) + "\n"
    text += str(i * 24) + "\n"
    text += (i * 12) * "#"
    return text

for i in range(20):
    label = str(i).zfill(4) + "." + ("ihatezoom" * i)
    f = pathlib.Path(label)
    out = data_folder / f
    out.write_text(make_text(i))


import pandas as pd
import numpy as np
from matplotlib import pyplot as plt
from pylab import *
mpl.rcParams['font.sans-serif'] = ['SimHei']

res = {}

for i in range(1, 16): # 统计15天的新增人数
    fileNameStr = './202012' + str(i).zfill(2) + '.csv'  # 产生文件名进行读取
    df = pd.read_csv(fileNameStr, encoding='utf-8')

    df['increase'] = df['increase'].astype(np.int)

    for idx in range(len(df)):  # 遍历所有国家
        if df['country'][idx] in res.keys():
            res[df['country'][idx]] = res[df['country'][idx]] + df['increase'][idx]
        else:
            res[df['country'][idx]] = df['increase'][idx]

lst = sorted(res.items(), key=lambda x:x[1], reverse=True)  # 按新增人数进行排序
country = []
increase = []
for i in range(10):  # 取出前10的国家
    country.append(lst[i][0])
    increase.append(lst[i][1])

plt.title("20201201~20201215 新冠病毒新增人数国家TOP10")
plt.bar(country, increase, label='increase')
plt.legend()
plt.show()


import pathlib
import numpy as np
import h5py
import cv2
import argparse

def load_episodes(directory, capacity=None):
    # The returned directory from filenames to episodes is guaranteed to be in
    # temporally sorted order.
    filenames = sorted(directory.glob('*.npz'))
    if capacity:
        num_steps = 0
        num_episodes = 0
        for filename in reversed(filenames):
            length = int(str(filename).split('-')[-1][:-4])
            num_steps += length
            num_episodes += 1
            if num_steps >= capacity:
                break
        filenames = filenames[-num_episodes:]
    episodes = {}
    for filename in filenames:
        try:
            with filename.open('rb') as f:
                episode = np.load(f)
                episode = {k: episode[k] for k in episode.keys()}
                # Conversion for older versions of npz files.
                if 'is_terminal' not in episode:
                    episode['is_terminal'] = episode['discount'] == 0.
        except Exception as e:
            print(f'Could not load episode {str(filename)}: {e}')
            continue
        episodes[str(filename)] = episode
    return episodes

def main():
    # Include argument parser
    parser = argparse.ArgumentParser(description='Convert npz files to hdf5.')
    parser.add_argument('--input_dir', type=str, required=True,
                        help='Path to input files')
    parser.add_argument('--output_dir', type=str, required=True,
                        help='Path to output files')
    args = parser.parse_args()

    step_type = np.ones(501)
    step_type[0] = 0
    step_type[500] = 2

    output = {}
    episodes = load_episodes(pathlib.Path(args.input_dir))
    episodes = list(episodes.values())

    actions = [e['action'] for e in episodes]
    discounts = [e['discount'] for e in episodes]
    observations = []
    for e in episodes:
        resized_images = np.empty((501, 84, 84, 3), dtype=e['image'].dtype)
        for (k, i) in enumerate(e['image']):
            resized_images[k] = cv2.resize(i, dsize=(84, 84), interpolation=cv2.INTER_CUBIC)
        observations.append(resized_images.transpose(0, 3, 1, 2))
    rewards = [e['reward'] for e in episodes]
    step_types = [step_type for _ in episodes]

    output['action'] = np.concatenate(actions)
    output['discount'] = np.concatenate(discounts)
    output['observation'] = np.concatenate(observations)
    output['reward'] = np.concatenate(rewards)
    output['step_type'] = np.concatenate(step_types)

    out_dir = pathlib.Path(args.output_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    with h5py.File(out_dir / 'data.hdf5', 'w') as shard_file:
        for k, v in output.items():
            shard_file.create_dataset(k, data=v, compression='gzip')

if __name__ == '__main__':
    main()


import json

from django.urls import reverse
from rest_framework import status
from rest_framework.test import APITestCase, APIClient
from django.contrib.auth import get_user_model
from rest_framework_simplejwt.tokens import RefreshToken

User = get_user_model()

AUTHOR = 'author'
EXECUTOR = 'executor'
AUTHOR_EMAIL = 'author@gmail.com'
NEW_AUTHOR_EMAIL = '123@gmail.com'
EXECUTOR_EMAIL = 'executor@gmail.com'
AUTHOR_ROLE = 'author'
EXECUTOR_ROLE = 'executor'
START_BALANCE = 500
NEW_BALANCE = 1
USERS_LIST_URL = reverse('users-list')

RESPOND_NEW_DATA = {'author': 2, 'task': 1}

class TaskModelTest(APITestCase):
    def setUp(self):
        self.admin = User.objects.create_superuser(
            username='admin',
            email='admin@gmail.com',
            balance=START_BALANCE,
            freeze_balance=START_BALANCE,
            is_staff=True,
            is_superuser=True)
        self.author = User.objects.create_user(
            username=AUTHOR,
            email=AUTHOR_EMAIL,
            balance=START_BALANCE,
            freeze_balance=START_BALANCE,
            role=AUTHOR_ROLE)
        self.executor = User.objects.create_user(
            username=EXECUTOR,
            email=EXECUTOR_EMAIL,
            balance=START_BALANCE,
            freeze_balance=START_BALANCE,
            role=EXECUTOR_ROLE)
        self.ADMIN_DETAIL_URL = reverse('users-detail', args=[self.admin.id])
        self.AUTHOR_ADD_BALANCE_URL = reverse('users-balance')
        self.USER_CHANGE_DATA_URL = reverse('users-me')
        self.AUTHOR_DETAIL_URL = reverse('users-detail', args=[self.author.id])
        self.EXECUTOR_DETAIL_URL = reverse('users-detail', args=[self.executor.id])
        self.executor_client = APIClient()
        self.executor_client.force_authenticate(user=self.executor)
        self.admin_client = APIClient()
        self.admin_client.force_authenticate(user=self.admin)
        self.token = RefreshToken.for_user(self.author)
        self.author_token = self.token.access_token
        self.auth_client = APIClient()
        self.auth_client.credentials(HTTP_AUTHORIZATION=f'Bearer {self.author_token}')

    def test_admin_get_users_list(self):
        response = self.admin_client.get(USERS_LIST_URL)
        self.assertEqual(response.status_code, status.HTTP_200_OK)
        self.assertEqual(len(response.data['results']), 3)

    def test_not_admin_cant_get_users_list(self):
        response = self.auth_client.get(USERS_LIST_URL)
        self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN)

    def test_admin_can_update_user_data(self):
        balance_before = self.author.balance
        response = self.admin_client.patch(
            self.AUTHOR_DETAIL_URL,
            data=json.dumps({'balance': NEW_BALANCE}),
            content_type='application/json')
        self.author.refresh_from_db()
        balance_after = self.author.balance
        self.assertEqual(response.status_code, status.HTTP_200_OK)
        self.assertEqual(balance_before, START_BALANCE)
        self.assertEqual(balance_after, NEW_BALANCE)

    def test_balance_url_add_money_to_balance(self):
        balance_before = self.author.balance
        response = self.auth_client.patch(
            self.AUTHOR_ADD_BALANCE_URL,
            data=json.dumps({'balance': NEW_BALANCE}),
            content_type='application/json')
        self.author.refresh_from_db()
        balance_after = self.author.balance
        self.assertEqual(response.status_code, status.HTTP_200_OK)
        self.assertEqual(balance_before, balance_after-NEW_BALANCE)

    def test_not_admin_cant_update_user_data(self):
        email_before = self.author.email
        response = self.auth_client.patch(
            self.AUTHOR_DETAIL_URL,
            data=json.dumps({'email': NEW_AUTHOR_EMAIL}),
            content_type='application/json')
        self.author.refresh_from_db()
        email_after = self.author.email
        self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN)
        self.assertEqual(email_before, email_after)

    def test_user_can_update_own_data(self):
        email_before = self.author.email
        response = self.auth_client.patch(
            self.USER_CHANGE_DATA_URL,
            data=json.dumps({'email': NEW_AUTHOR_EMAIL}),
            content_type='application/json')
        self.author.refresh_from_db()
        email_after = self.author.email
        self.assertEqual(response.status_code, status.HTTP_200_OK)
        self.assertEqual(email_before, AUTHOR_EMAIL)
        self.assertEqual(email_after, NEW_AUTHOR_EMAIL)


import os
from cffi import FFI
from OpenSSL.SSL import Context as SSLContext, _ffi, _lib as lib

from utils import OutputGrabber

ffi = FFI()
NULL = ffi.NULL
ffi.cdef(
    "int SSL_CTX_set_client_cert_engine(void *ctx, void *e);"
    "int ENGINE_set_default(void *e, unsigned int flags);"
)
libcrypto = ffi.dlopen("libcrypto-1_1.dll")
libssl = ffi.dlopen("libssl-1_1.dll")

class ENGINE_DEFAULT:
    ALL = 0xFFFF

class CAPI_LIST_DISP_FMT:
    SUMMARY = 1
    FRIENDLY_NAME = 2
    FULL = 4
    PEM = 8
    XXX = 16
    PRIV_KEY_INFO = 32

class SSLEngine(object):
    def __init__(self, id: str | FFI.CData) -> None:
        if isinstance(id, str):
            try:
                eng = SSLEngine.load_by_id(id)
            except Exception:
                eng = SSLEngine.load_dynamic(id)
            ptr = eng.ptr
        elif isinstance(id, SSLEngine):
            ptr = id.ptr
        else:
            ptr = id
        self.ptr = ptr

    def init(self):
        if not lib.ENGINE_init(self.ptr):
            self.__exit__()
            raise Exception("Could not initialize engine")

    def free(self):
        lib.ENGINE_free(self.ptr)

    def __enter__(self):
        self.init()
        return self

    def __exit__(self, type, value, traceback):
        self.free()

    def set_default(self, flags: int = ENGINE_DEFAULT.ALL):
        if not libcrypto.ENGINE_set_default(self.ptr, flags):
            self.free()
            raise Exception(
                "Not able to set engine as default for all flags:%s" % flags
            )

    def ctrl_cmd_string(
        self,
        cmd: str,
        value: str | None = None,
        optional: bool = False,
        capture: bool = False,
    ) -> None | bytes:
        io: None | OutputGrabber = None
        if capture:
            io = OutputGrabber(threaded=True)
            io.start()
        if not lib.ENGINE_ctrl_cmd_string(
            self.ptr,
            cmd.encode("utf-8"),
            NULL if value == None else value.encode("utf-8"),
            1 if optional else 0,
        ):
            if capture:
                io.stop()
            raise Exception(
                "Error with engine string control command: %s%s"
                % (cmd, "" if value == None else ":" + value)
            )
        if capture:
            io.stop()
            return io.captured

    def load_by_id(id: str):
        if not id:
            raise ValueError("Id value must be provided")
        lib.ENGINE_load_builtin_engines()
        ptr = lib.ENGINE_by_id(id.encode())
        if ptr == NULL:
            raise ValueError("Could not load the {0} engine by id".format(id))
        return SSLEngine(ptr)

    def load_dynamic(
        id: str,
        path: str = None,
        search_path: str = None,
        check_version: bool = True,
    ):

        if not id:
            raise ValueError("Id value must be provided")

        dyn = SSLEngine.load_by_id("dynamic")
        dyn.ctrl_cmd_string("ID", id)

        if path:
            dyn.ctrl_cmd_string("SO_PATH", path)

        dyn.ctrl_cmd_string("LIST_ADD", "1")

        if not check_version:
            dyn.ctrl_cmd_string("NO_VCHECK", "1")

        if search_path == None and path == None and "OPENSSL_ENGINES" in os.environ:
            search_path = os.environ ["OPENSSL_ENGINES"]

        if search_path:
            dyn.ctrl_cmd_string("DIR_LOAD", "2")
            dyn.ctrl_cmd_string("DIR_ADD", search_path)

        dyn.ctrl_cmd_string("LOAD")
        return dyn

class CAPIEngine(SSLEngine):
    def __init__(self, src: FFI.CData | str | SSLEngine | None = None) -> None:
        super().__init__("capi" if src == None else src)

    def set_store(self, name: str):
        self.ctrl_cmd_string("store_name", name)

    def list_certs(
        self, store: str | None = None, format: int | None = None
    ) -> list[bytes]:
        if format:
            self.ctrl_cmd_string("list_options", str(format))
        if store:
            self.set_store(store)
        return [
            cert.split(sep=b"\n", maxsplit=1)[1]
            for cert in self.ctrl_cmd_string("list_certs", capture=True)
            .strip(b"\n")
            .split(b"\nCertificate ")
        ]

def set_client_cert_engine(self: SSLContext, engine: FFI.CData | SSLEngine):
    if not libssl.SSL_CTX_set_client_cert_engine(
        self._context, engine.ptr if isinstance(engine, SSLEngine) else engine
    ):
        raise Exception("Was not able to set client cert engine")

SSLContext.set_client_cert_engine = set_client_cert_engine


#!/usr/bin/env python3

"""
Implementation of R3PTAR
"""

import logging
import signal
import sys
from ev3dev2.motor import OUTPUT_A, OUTPUT_B, OUTPUT_C, OUTPUT_D, MediumMotor, LargeMotor
from ev3dev2.sensor.lego import InfraredSensor
from ev3dev2.sound import Sound
from threading import Thread, Event
from time import sleep

log = logging.getLogger(__name__)

class MonitorRemoteControl(Thread):
    """
    A thread to monitor R3PTAR's InfraredSensor and process signals
    from the remote control
    """

    def __init__(self, parent):
        Thread.__init__(self)
        self.parent = parent
        self.shutdown_event = Event()

    def __str__(self):
        return "MonitorRemoteControl"

    def run(self):
        STRIKE_SPEED_PCT = 40

        while True:

            if self.shutdown_event.is_set():
                log.info('%s: shutdown_event is set' % self)
                break

            #log.info("proximity: %s" % self.parent.remote.proximity)
            if self.parent.remote.proximity < 30:
                self.parent.speaker.play('snake-hiss.wav', Sound.PLAY_NO_WAIT_FOR_COMPLETE)
                self.parent.strike_motor.on_for_seconds(speed=STRIKE_SPEED_PCT, seconds=0.5)
                self.parent.strike_motor.on_for_seconds(speed=(STRIKE_SPEED_PCT * -1), seconds=0.5)

            self.parent.remote.process()
            sleep(0.01)

class R3PTAR(object):

    def __init__(self,
                 drive_motor_port=OUTPUT_B,
                 strike_motor_port=OUTPUT_D,
                 steer_motor_port=OUTPUT_A,
                 drive_speed_pct=60):

        self.drive_motor = LargeMotor(drive_motor_port)
        self.strike_motor = LargeMotor(strike_motor_port)
        self.steer_motor = MediumMotor(steer_motor_port)
        self.speaker = Sound()
        STEER_SPEED_PCT = 30

        self.remote = InfraredSensor()
        self.remote.on_channel1_top_left = self.make_move(self.drive_motor, drive_speed_pct)
        self.remote.on_channel1_bottom_left = self.make_move(self.drive_motor, drive_speed_pct * -1)
        self.remote.on_channel1_top_right = self.make_move(self.steer_motor, STEER_SPEED_PCT)
        self.remote.on_channel1_bottom_right = self.make_move(self.steer_motor, STEER_SPEED_PCT * -1)

        self.shutdown_event = Event()
        self.mrc = MonitorRemoteControl(self)

        # Register our signal_term_handler() to be called if the user sends
        # a 'kill' to this process or does a Ctrl-C from the command line
        signal.signal(signal.SIGTERM, self.signal_term_handler)
        signal.signal(signal.SIGINT, self.signal_int_handler)

    def make_move(self, motor, speed):
        def move(state):
            if state:
                motor.on(speed)
            else:
                motor.stop()
        return move

    def shutdown_robot(self):

        if self.shutdown_event.is_set():
            return

        self.shutdown_event.set()
        log.info('shutting down')
        self.mrc.shutdown_event.set()

        self.remote.on_channel1_top_left = None
        self.remote.on_channel1_bottom_left = None
        self.remote.on_channel1_top_right = None
        self.remote.on_channel1_bottom_right = None

        self.drive_motor.off(brake=False)
        self.strike_motor.off(brake=False)
        self.steer_motor.off(brake=False)

        self.mrc.join()

    def signal_term_handler(self, signal, frame):
        log.info('Caught SIGTERM')
        self.shutdown_robot()

    def signal_int_handler(self, signal, frame):
        log.info('Caught SIGINT')
        self.shutdown_robot()

    def main(self):
        self.mrc.start()
        self.shutdown_event.wait()

if __name__ == '__main__':

    # Change level to logging.DEBUG for more details
    logging.basicConfig(level=logging.INFO,
                        format="%(asctime)s %(levelname)5s %(filename)s: %(message)s")
    log = logging.getLogger(__name__)

    # Color the errors and warnings in red
    logging.addLevelName(logging.ERROR, "\033[91m  %s\033[0m" % logging.getLevelName(logging.ERROR))
    logging.addLevelName(logging.WARNING, "\033[91m%s\033[0m" % logging.getLevelName(logging.WARNING))

    log.info("Starting R3PTAR")
    snake = R3PTAR()
    snake.main()
    log.info("Exiting R3PTAR")


#!/usr/bin/env python

import numpy as np
import healpy as hp
import pylab
import matplotlib.pyplot as plt
import time
import mocklc
import matplotlib
import sepmat
import gpkernel
import scipy
import emcee
import sys
import time

Ns=2000

np.random.seed(17)

#set geometry
inc=45.0/180.0*np.pi
Thetaeq=np.pi
zeta=60.0/180.0*np.pi
Pspin=23.9344699/24.0 #Pspin: a sidereal day
wspin=2*np.pi/Pspin
Porb=365.242190402
worb=2*np.pi/Porb
Ni=1024
obst=np.linspace(0.0,Porb,Ni)

# test map
nside=16
npix=hp.nside2npix(nside)
mmap=hp.read_map("/home/kawahara/exomap/sot/data/mockalbedo16.fits")
mask=(mmap>0.0)
mmap[mask]=1.0
M=len(mmap)

#generating light curve
Thetav=worb*obst
Phiv=np.mod(wspin*obst,2*np.pi)
WI,WV=mocklc.comp_weight(nside,zeta,inc,Thetaeq,Thetav,Phiv)
W=WV[:,:]*WI[:,:]
lc=np.dot(W,mmap)

sigma=np.mean(lc)*0.1
noise=np.random.normal(0.0,sigma,len(lc))
lc=lc+noise

## RBF kernel
nside=16
npix=hp.nside2npix(nside)
sep=sepmat.calc_sepmatrix(nside)

## optimization

tag="RBFobl"
## spin and hyperparameter MCMC sampling using emcee
def log_prior(theta):
    p_zeta,p_Thetaeq,p_gamma,p_alpha=theta
    if 0.0 < p_zeta < np.pi and 0.0 < p_Thetaeq < 2*np.pi and 1.e-10 < p_gamma < np.pi/3.0 and 1.e-10 < p_alpha:
        return np.log(np.sin(p_zeta)/p_alpha/p_gamma)
    return -np.inf

def log_likelihood(theta, d, covd):
    p_zeta,p_Thetaeq,p_gamma,p_alpha=theta
    WI,WV=mocklc.comp_weight(nside,p_zeta,inc,p_Thetaeq,Thetav,Phiv)
    Wp=WV[:,:]*WI[:,:]
    #KS=p_alpha*gpkernel.Matern32(sep,p_gamma)
    KS=p_alpha*gpkernel.RBF(sep,p_gamma)
    Cov = covd + Wp@KS@Wp.T
    sign,logdet=np.linalg.slogdet(Cov)
    Pi_d=scipy.linalg.solve(Cov,d,assume_a="pos")
    prop = -0.5*logdet-0.5*d@Pi_d #-0.5*np.shape(cov)[0]*np.log(2.0*np.pi)
    return prop

def log_probability(theta, d, covd):
    lp = log_prior(theta)
    if not np.isfinite(lp):
        return -np.inf
    return lp + log_likelihood(theta, d, covd)

gam0=0.29298260376811
alpha0=sigma**2*0.774263682681127
pos = np.array([zeta,Thetaeq,gam0,alpha0])+ 1e-4 * np.random.randn(16, 4)
nwalkers, ndim = pos.shape

#Assumming we know the data covariance
covd=sigma**2*np.eye(Ni)

sampler = emcee.EnsembleSampler(nwalkers, ndim, log_probability, args=(lc, covd))
sampler.run_mcmc(pos, Ns, progress=True);

flat_samples = sampler.get_chain(discard=100, thin=15, flat=True)
#samples = sampler.get_chain()
#print(samples)
labels=["zeta","Thetaeq","gamma","alpha"]

inputgeo=[inc,Thetaeq,zeta,Pspin,Porb,obst]
np.savez("flat_sample"+tag,flat_samples,W,lc,inputgeo)

import corner
fig = corner.corner(flat_samples, labels=labels, truths=[zeta,Thetaeq,None,None])
plt.savefig("corner_"+tag+".png")
plt.savefig("corner_"+tag+".pdf")
plt.show()


#!/usr/bin/env python3
import argparse
import os
import re
import dataclasses
from dataclasses import dataclass
from pathlib import Path
from typing import Optional, List

"""
Supports following cases:

1. Master version x.y.z needs to be bumped to x.y.z when preparing for official release:

    git checkout cluster-test
    git merge master
    # version = x.y.z

    version_bumper.py
    # version = x.y.z+1

    version_bumper.py --part=minor
    # version = x.y+1.0

    version_bumper.py --part=major
    # version = x+1.0.0

2. Master version x.y.z needs to be bumped to x.y.z-mr-1 when making dev release from feature branch:

    git co 123-my-branch
    # version = x.y.z

    version_bumper.py --mr 123
    # version = x.y.z-123-1

 And then another call should just bump the dev-version:
    version_bumper.py --mr 123
    # version = x.y.z-123-2

"""

@dataclass
class Version:
    major: int
    minor: int
    patch: int
    mr: int   # merge request id
    dev: int  # sequentially increasing number

    def __str__(self):
        mr = f"-{self.mr}" if self.mr > 0 else ''
        dev = f"-{self.dev}" if self.dev > 0 else ''
        return f'{self.major}.{self.minor}.{self.patch}{mr}{dev}'

    def bump(self, part: str):
        self.__dict__[part] += 1
        if part == 'major':
            self.minor = self.patch = 0
        if part == 'minor':
            self.patch = 0

    def clone(self) -> 'Version':
        return dataclasses.replace(self)

def read_current_version(filepath: Path) -> Version:
    for line in filepath.read_text().splitlines():
        ver = parse_version(line)
        if ver is not None:
            return ver
    raise RuntimeError('version could not be parsed from ' + str(filepath))

# match X.Y.Z or X.Y.Z-W
# broken down at https://regex101.com/r/IAccOs/3
main_regex = r'TAG\s*\?=\s*?(?P<major>\d+)\.(?P<minor>\d+)\.(?P<patch>\d+)[\-\.]?(?P<details>[\-\w]+)?'
main_pattern = re.compile(main_regex)

def parse_version(line: str) -> Optional[Version]:
    match = main_pattern.match(line)
    if not match:
        return None
    ver = Version(major=int(match.group('major')),
                  minor=int(match.group('minor')),
                  patch=int(match.group('patch')),
                  mr=0,
                  dev=0)
    details = match.group('details')
    if details is not None:
        parse_details(details, ver)
    return ver

# match X-Y
# broken down at https://regex101.com/r/jtlQ54/3
details_regex = r'(?P<mr>\d+)[\-](?P<dev>\d+)'
details_pattern = re.compile(details_regex)

def parse_details(details: str, ver: Version):
    details_match = details_pattern.match(details)
    if details_match:
        ver.mr = int(details_match.group('mr'))
        ver.dev = int(details_match.group('dev'))

def replace_in_files(curr_ver: Version, new_ver: Version, files: List[Path]):
    for path in files:
        replace_in_file(path, curr_ver, new_ver)

def replace_in_file(filepath: Path, curr_ver: Version, new_ver: Version):
    content = filepath.read_text()
    new_content = content.replace(str(curr_ver), str(new_ver))
    if content != new_content:
        filepath.write_text(new_content)
        print(f'Version bumped {curr_ver} -> {new_ver} in {filepath}')
    else:
        raise RuntimeError(f'Version "{curr_ver}" not found in {filepath}')

def project_root() -> Path:
    """Return Racetrack root dir"""
    return Path(os.path.abspath(__file__)).parent.parent.absolute()

def bump_version_in_files(version_path: Path, _args, files: List[Path], prod_files: List[Path]):
    orig_version = read_current_version(version_path)

    if _args.current:
        print(orig_version)
        return

    new_version = orig_version.clone()
    if _args.mr and int(_args.mr) != 0:
        new_version.mr = int(_args.mr)

    if new_version.mr != 0:
        new_version.bump('dev')
    else:
        new_version.bump(_args.part)
        files += prod_files
    replace_in_files(orig_version, new_version, files)

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--current', action='store_true', help='print current version')
    parser.add_argument('--mr', help='set merge request number')
    parser.add_argument('--part', help='defines which part to bump: major, minor, patch, dev', default="patch")

    files_with_version = [
        project_root() / 'Makefile',
    ]
    # files bumped in official (non-dev) releases only
    prod_files_with_version = [
        project_root() / 'racetrack_client/racetrack_client/__init__.py',
    ]

    args = parser.parse_args()
    path = project_root() / 'Makefile'
    bump_version_in_files(path, args, files_with_version, prod_files_with_version)


#!/usr/bin/env python2

import argparse
import ast
import json
import logging
import os
from collections import namedtuple
import tqdm

import sys
sys.path.append('.')
print(sys.path)

from srcseq.astunparser import Unparser, WriterBase

def file_tqdm(fobj):
    return tqdm(fobj, total=get_number_of_lines(fobj))

SrcASTToken = namedtuple("SrcASTToken", "text type lineno col_offset")
logging.basicConfig(level=logging.INFO)

class MyListFile(list, WriterBase):
    def write(self, text, type=None, node=None):
        text = text.strip()
        lineno = node and node.lineno
        col_offset = node and node.col_offset
        if len(text) > 0:
            # write `Str` as it is. `Num` will be kept as a string.
            text = eval(text) if type == "Str" else text
            self.append(SrcASTToken(text, type, lineno, col_offset))

    def flush(self):
        pass

def my_tokenize(code_str):
    t = ast.parse(code_str)
    lst = MyListFile()
    Unparser(t, lst)
    return lst

def main():
    parser = argparse.ArgumentParser(
        description="Generate datapoints from source code",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument("--files_path", "-f", required=True,
        help="Filepath with the filenames to be parsed")
    parser.add_argument("--save", "-o", default="/tmp/dps.jsonl",
        help="Filepath with the output dps")
    parser.add_argument("--base_dir", "-b",
        help="Base dir to append for the fps."
        " If not given, use the dir of `--files_path`.")
    args = parser.parse_args()
    args.base_dir = args.base_dir or os.path.dirname(args.files_path)
    if os.path.exists(args.save):
        os.remove(args.save)

    num_dps = 0
    logging.info("Loading files from: {}".format(args.base_dir))
    with open(args.files_path, "r") as fin, open(args.save, "w") as fout:
        for i_line, line in enumerate(file_tqdm(fin)):
            rel_src_fp = line.strip()
            abs_src_fp = os.path.join(args.base_dir, rel_src_fp)
            try:
                values, types_, linenos, col_offsets = zip(*my_tokenize(open(abs_src_fp).read()))
                if len(values) > 1:
                    json.dump({
                        'rel_src_fp': rel_src_fp,
                        'values': values,
                        'types': types_,
                        'linenos': linenos,
                        'col_offsets': col_offsets,
                    }, fp=fout)
                    fout.write("\n")
                    num_dps += 1
                else:
                    # logging.info("In processing {}-th file `{}`: empty token list.".format(i_line, rel_src_fp))
                    pass
            except Exception as e:
                logging.warning("In processing {}-th file `{}`:\n\t{}".format(i_line, rel_src_fp, e))
                continue
    logging.info("Wrote {} datapoints to {}".format(num_dps, args.save))

if __name__ == "__main__":
    main()