from matchmaking import sheets, match, mmdiscord
from discord.ext import commands
import discord.utils
from settings import *
import random
import asyncio
from country_list import countries_for_language
from matchmaking.mmdiscord import founder, admin, moder, member

class Matchmaking(commands.Cog):

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

    @commands.command(name='class', brief='Changes your classes. Use !mm class main_class secondary_class')
    @member()
    async def classes(self, ctx, main=None, secondary=None, player=None):
        if (not main) or (not secondary):
            await ctx.send('Please add both main and secondary classes ', delete_after=5)
            await ctx.message.delete()
            return

        if not (main in role_fix_names):
            await ctx.send('No such class %s found' % main, delete_after=5)
            await ctx.message.delete()
            return

        if not (secondary in role_fix_names):
            await ctx.send('No such class %s found' % secondary, delete_after=5)
            await ctx.message.delete()
            return

        if sheets.role_fix_names(main) == sheets.role_fix_names(secondary):
            await ctx.send('Main and secondary classes should be different', delete_after=5)
            await ctx.message.delete()
            return

        if player:
            admin_role = discord.utils.find(lambda x: x.name == sheet_edit_perms, ctx.message.guild.roles)
            if admin_role in ctx.author.roles:

                member = discord.utils.find(lambda x: x.display_name.lstrip(' ◦◊❃⌇୭૬৸৶৯ৡঙঐ') == player,
                                            ctx.guild.members)

                if not member:
                    await ctx.send('No such discord user %s found' % player, delete_after=5)
                    await ctx.message.delete()
                    return

            else:
                await ctx.send('No permissions', delete_after=5)
                await ctx.message.delete()
                return
        else:
            member = ctx.author

        found_player = sheets.find_player(member)
        if not found_player:
            await ctx.send('No player %s found' % member.display_name.lstrip(' ◦◊❃⌇୭૬৸৶৯ৡঙঐ'), delete_after=5)
            await ctx.message.delete()
            return

        sheets.change_classes(member, main, secondary)
        await ctx.send('Classes of player %s have been changed to Main: %s and Secondary: %s' %
                       (member.name, sheets.role_fix_names(main), sheets.role_fix_names(secondary)),
                       delete_after=10)

        await ctx.message.delete()

    @commands.command(brief='Changes your clan. Use !mm clan clan_name')
    @admin()
    async def clan(self, ctx, clan, player=None):
        if player:
            admin_role = discord.utils.find(lambda x: x.name == sheet_edit_perms, ctx.message.guild.roles)
            if admin_role in ctx.author.roles:

                member = discord.utils.find(lambda x: x.display_name.lstrip(' ◦◊❃⌇୭૬৸৶৯ৡঙঐ') == player,
                                            ctx.guild.members)

                if not member:
                    await ctx.send('No such discord user %s found' % player, delete_after=5)
                    await ctx.message.delete()
                    return

            else:
                await ctx.send('No permissions', delete_after=5)
                await ctx.message.delete()
                return
        else:
            member = ctx.author

        found_player = sheets.find_player(member)
        if not found_player:
            await ctx.send('No player %s found' % member.display_name.lstrip(' ◦◊❃⌇୭૬৸৶৯ৡঙঐ'), delete_after=5)
            await ctx.message.delete()
            return

        if clan in clan_list or clan in no_clan:
            sheets.change_clan(member, clan)
            await ctx.send('Changed the clan of player %s to %s' % (member.name, clan), delete_after=10)
        else:
            await ctx.send(('No clan %s in the clan list. \nCurrent clan list = ' % clan) + str(clan_list) +
                           '. \nContact Matchmaking admininstration to register a new clan', delete_after=10)

        await ctx.message.delete()

    @commands.command()
    @admin()
    async def country(self, ctx, country, player=None):
        if player:
            admin_role = discord.utils.find(lambda x: x.name == sheet_edit_perms, ctx.message.guild.roles)
            if admin_role in ctx.author.roles:

                member = discord.utils.find(lambda x: x.display_name.lstrip(' ◦◊❃⌇୭૬৸৶৯ৡঙঐ') == player,
                                            ctx.guild.members)

                if not member:
                    await ctx.send('No such discord user %s found' % player, delete_after=5)
                    await ctx.message.delete()
                    return

            else:
                await ctx.send('No permissions', delete_after=5)
                await ctx.message.delete()
                return
        else:
            member = ctx.author

        found_player = sheets.find_player(member)
        if not found_player:
            await ctx.send('No player %s found' % member.display_name.lstrip(' ◦◊❃⌇୭૬৸৶৯ৡঙঐ'), delete_after=5)
            await ctx.message.delete()
            return

        country_dict = dict(countries_for_language('en'))

        if (not (country in country_dict.values())) and (not (country in country_dict.keys())):
            await ctx.send('No country %s found in country list' % country, delete_after=5)
            await ctx.message.delete()
            return

        if country in country_dict.keys():
            country = country_dict[country]

        sheets.change_country(member, country)
        await ctx.send('Country of player %s changed to %s' % (member.name, country), delete_after=20)
        await ctx.message.delete()

    @commands.command(brief='Use !mm info for your own stats or !mm info [player1] [player2]')
    @member()
    async def info(self, ctx, *args):
        print(str(ctx.message.author), ': ', str(ctx.message.content))
        author = ctx.message.author
        if len(args) == 0:
            args = [author.display_name.lstrip(' ◦◊❃⌇୭૬৸৶৯ৡঙঐ')]
        if len(args) == 1 and args[0] == 'voice':
            if ctx.author.voice:
                member_list = []
                not_found_players_list = []

                for member in ctx.author.voice.channel.members:
                    found_player = sheets.find_player(member)
                    if found_player:
                        member_list.append(found_player)
                    else:
                        not_found_players_list.append(member)

                member_list.sort(key=lambda x: x.mmr, reverse=True)

                for member in member_list:
                    embed = mmdiscord.form_player_embed(member)
                    await ctx.send(embed=embed, delete_after=600)

                not_found_players_message = 'Players '
                if not_found_players_list:
                    for not_found_player in not_found_players_list:

                        not_found_players_message += not_found_player.name
                        not_found_players_message += ', '

                    not_found_players_message = not_found_players_message[:-2]
                    not_found_players_message += ' not found'

                    await ctx.send(not_found_players_message, delete_after=600)

            else:
                await ctx.send("You're not in a voice channel", delete_after=10)

        else:
            for member_name in args:
                member = discord.utils.find(lambda x: x.display_name.lstrip(' ◦◊❃⌇୭૬৸৶৯ৡঙঐ') ==
                                            member_name, ctx.guild.members)
                if member:
                    founded_player = sheets.find_player(member)
                    if founded_player:
                        embed = mmdiscord.form_player_embed(founded_player)
                        await ctx.send(embed=embed, delete_after=600)
                    else:
                        await ctx.send(('No player %s found' % member.display_name.lstrip(' ◦◊❃⌇୭૬৸৶৯ৡঙঐ')),
                                       delete_after=10)
                else:
                    message = 'No discord user %s found' % member_name
                    await ctx.send(message, delete_after=10)
        await ctx.message.delete()

    @commands.command(brief='Use !mm reg (main class) (secondary class) (clan name)')
    @member()
    async def reg(self, ctx, *args):
        print(str(ctx.message.author), ': ', str(ctx.message.content))
        member = ctx.message.author
        correct_request = True

        if not len(args) == 1:
            if len(args) == 2 or len(args) == 3:

                if not args[0] in role_fix_names:
                    await ctx.send('No such role ' + args[0], delete_after=10)
                    correct_request = False

                elif not args[1] in role_fix_names:
                    await ctx.send('No such role ' + args[1], delete_after=10)
                    correct_request = False

                elif sheets.role_fix_names(args[0]) == sheets.role_fix_names(args[1]):
                    await ctx.send('Classes should be different', delete_after=10)
                    correct_request = False

            if len(args) == 3 and correct_request:
                if args[2] not in clan_list:
                    await ctx.send('No such clan ' + args[2] + ' in the list', delete_after=10)
                    correct_request = False
            else:
                if len(args) > 3:
                    await ctx.send('Too many arguments', delete_after=10)
                    correct_request = False
        else:
            await ctx.send('2 Classes required!', delete_after=10)
            correct_request = False

        if correct_request:
            if not sheets.find_player(member):
                sheets.add_player(member, *args)
                await ctx.send('Player ' + member.display_name + ' added')
            else:
                await ctx.send('Player ' + member.display_name + ' already exists, use !mm class and !mm clan',
                               delete_after=10)
        await ctx.message.delete()

    @commands.command(brief='Starts Matchmaking, 12 players needed in a voice channel')
    @member()
    async def start(self, ctx):
        if self.mm_starts:
            return
        self.mm_starts = True

        print(str(ctx.message.author), ': ', str(ctx.message.content))

        if not ctx.channel.name == team_building_text_channel_name:
            print('Wrong text channel')
            return

        if not ctx.author.voice.channel.name == team_building_voice_channel_name:
            print('Wrong voice channel')
            return

        if ctx.author.voice:
            voice_members = ctx.author.voice.channel.members
            players_amount = len(voice_members)

            if players_amount < 12:
                await ctx.send('Not enough players, need ' + str(12 - players_amount) + ' more', delete_after=10)
            else:

                removed_members = []
                while (len(voice_members) % 12) != 0:

                    removed_member = random.choice(voice_members)
                    voice_members.remove(removed_member)
                    removed_members.append(removed_member)

                await ctx.send(random.choice(waiting_phrases), delete_after=5)
                users = voice_members
                players = []
                player_count = 0

                for mm_player in users:
                    player_count += 1
                    player_to_add = sheets.find_player(mm_player)

                    if player_count % 10 == 0:
                        await ctx.send(random.choice(waiting_phrases), delete_after=5)

                    if player_to_add:
                        players.append(player_to_add)

                    else:
                        sheets.add_player(mm_player)
                        players.append(sheets.find_player(mm_player))

                await ctx.send(random.choice(waiting_phrases), delete_after=10)

                players.sort(key=lambda x: x.mmr, reverse=True)

                team1_voice_channel = discord.utils.find(lambda x: x.name == team1_voice_channel_name,
                                                         ctx.channel.category.channels)
                team3_voice_channel = discord.utils.find(lambda x: x.name == team3_voice_channel_name,
                                                         ctx.channel.category.channels)
                team5_voice_channel = discord.utils.find(lambda x: x.name == team5_voice_channel_name,
                                                         ctx.channel.category.channels)
                team7_voice_channel = discord.utils.find(lambda x: x.name == team7_voice_channel_name,
                                                         ctx.channel.category.channels)

                while len(players) > 0:
                    game_players = []
                    while len(game_players) != 12:
                        game_players.append(players[0])
                        players.pop(0)

                    game = match.Match(game_players)
                    game.set_roles()
                    game.make_teams()

                    team1_players = []
                    team2_players = []

                    for match_player in game.team1:
                        team1_players.append(ctx.guild.get_member(int(match_player.discord_id)))
                    for match_player in game.team2:
                        team2_players.append(ctx.guild.get_member(int(match_player.discord_id)))

                    if not team1_voice_channel.members:
                        team2_voice_channel = discord.utils.find(lambda x: x.name == team2_voice_channel_name,
                                                                 ctx.channel.category.channels)
                        team1_founded_voice_channel = team1_voice_channel
                        team2_founded_voice_channel = team2_voice_channel
                    elif not team3_voice_channel.members:
                        team4_voice_channel = discord.utils.find(lambda x: x.name == team4_voice_channel_name,
                                                                 ctx.channel.category.channels)
                        team1_founded_voice_channel = team3_voice_channel
                        team2_founded_voice_channel = team4_voice_channel
                    elif not team5_voice_channel.members:
                        team6_voice_channel = discord.utils.find(lambda x: x.name == team6_voice_channel_name,
                                                                 ctx.channel.category.channels)
                        team1_founded_voice_channel = team5_voice_channel
                        team2_founded_voice_channel = team6_voice_channel
                    elif not team7_voice_channel.members:
                        team8_voice_channel = discord.utils.find(lambda x: x.name == team8_voice_channel_name,
                                                                 ctx.channel.category.channels)
                        team1_founded_voice_channel = team7_voice_channel
                        team2_founded_voice_channel = team8_voice_channel
                    else:
                        await ctx.send('No free voice channels found')
                        return

                    for member in team1_players:
                        await member.move_to(team1_founded_voice_channel)
    #                    await asyncio.sleep(0.3)

                    await ctx.send('Moving players...', delete_after=20)

                    for member in team2_players:
                        await member.move_to(team2_founded_voice_channel)
    #                    await asyncio.sleep(0.3)

                    if sort_teams_by == 'MMR':
                        game.sort_teams_by_mmr()

                    embed = mmdiscord.form_teams_embed(game.team1, game.team2)
                    await ctx.send(embed=embed)

        else:
            await ctx.send('You are not in a voice channel', delete_after=10)

        await ctx.message.delete()
        self.mm_starts = False

    @commands.command(brief='Updates MMR')
    @admin()
    async def update(self, ctx):
        print(str(ctx.message.author), ': ', str(ctx.message.content))
        print('Updating...')
        sheets.insert_formulas()
        await ctx.send('MMR updated!', delete_after=10)
        await ctx.message.delete()

    @commands.command(brief='Moves everyone to Team Building channel, only for modders')
    @moder()
    async def reset(self, ctx):

        print(str(ctx.message.author), ': ', str(ctx.message.content))
        await ctx.message.delete()

        if not (ctx.author.voice.channel.name in (team2_voice_channel_name, team1_voice_channel_name)):
            print('Wrong voice channel')
            return

        direct_channel = discord.utils.find(lambda x: x.name == team_building_voice_channel_name,
                                            ctx.channel.category.channels)
        team1_channel = discord.utils.find(lambda x: x.name == team1_voice_channel_name,
                                           ctx.channel.category.channels)
        team2_channel = discord.utils.find(lambda x: x.name == team2_voice_channel_name,
                                           ctx.channel.category.channels)

        for member1 in team1_channel.members:
            await member1.move_to(direct_channel)
            await asyncio.sleep(0.3)
        for member2 in team2_channel.members:
            await member2.move_to(direct_channel)
            await asyncio.sleep(0.3)

        await ctx.send('Game reset', delete_after=20)

    @commands.command(brief='Locks MMR', name='lockmmr')
    @founder()
    async def new_rating_system(self, ctx):
        print(str(ctx.message.author), ': ', str(ctx.message.content))
        print('Updating...')
        sheets.replace_formulas_with_values()
        await ctx.send('MMR locked')

def setup(bot):
    bot.add_cog(Matchmaking(bot))


# -*- coding: utf-8 -*-
"""
Created on Wed Aug  5 15:27:58 2020

@author: engelen
"""

import xarray as xr
import os
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1.inset_locator import inset_axes
import seaborn as sns

from itertools import product

def prepare_for_plot(ds):
    ds = ds.where(ds > -9998.)
    ds = ds.assign_coords(x = ds["x"]/100000)
    ds = ds.assign_coords(z = ds["z"]/10)
    ds = ds.reset_coords(drop=True)
    ds = ds.clip(min=0., max=35.)
    return(ds)

#%%Path management
outf = r"c:\Users\engelen\OneDrive - Stichting Deltares\PhD\Synth_Delta\results\2D_cross"

folders = [#r"g:\synthdelta\results_SA_zip\synth_SD_i018_m24_7080926",
           r"g:\synthdelta\results_SA_zip\synth_SD_i147_m24_7114854",
           r"g:\synthdelta\results_SA_zip\synth_SD_i043_m24_7065053",
           r"g:\synthdelta\results_SA_zip\synth_SD_i049_m24_7087160",
           r"g:\synthdelta\results_SA_zip\synth_SD_i123_m24_7114121",
           r"g:\synthdelta\results_SA_zip\synth_SD_i140_m24_7114849"]

files = ["results_006.nc",
         "results_006.nc",
         "results_006.nc",
         "results_001.nc",
         "results_006.nc"]

files = [os.path.join(folder, f) for folder, f in zip(folders, files)]

#%%Classification
types = [3, 5, 4, 1, 2]

#%%Reorder
files = [x for _,x in sorted(zip(types,files))]

#%%Read and select
ds_ls = [xr.open_dataset(f)["conc1"].isel(time=-1) for f in files]
ds_2d = [ds.isel(y=-1) for ds in ds_ls]

#%%Prepare for plot
ds_2d = [prepare_for_plot(ds) for ds in ds_2d]

#%%Plot
colors = [
 [0.171875, 0.48046875, 0.7109375],
 [0.5703125, 0.7734375, 0.87109375],
 [0.8671875, 0.9375, 0.8125],
 [0.99609375, 0.87109375, 0.6015625],
 [0.9609375, 0.5625, 0.32421875],
 [0.83984375, 0.09765625, 0.109375],
 [0.62890625, 0.09765625, 0.109375],
 # [0.421875, 0.09765625, 0.109375],
]

# levels = [0, 1, 5, 10, 15, 20, 25, 30, 35]
levels = [0, 1, 5, 10, 15, 20, 25, 35]

#http://xarray.pydata.org/en/stable/generated/xarray.DataArray.plot.contourf.html
#ds_2d[0].plot.contourf(levels = levels, colors = colors)

#%%Setup Plot
widths = [10, 10, 1]
widths = [10, 10]

dark = ".1"
light = ".95"

style_dict = {"axes.facecolor": dark,
               "axes.labelcolor" : dark,
               'text.color': light,
               'figure.facecolor': dark,
               'xtick.color': light,
               'ytick.color': light}

with sns.axes_style("dark", style_dict):
    with sns.plotting_context("paper", font_scale=0.7):
        fig = plt.figure(constrained_layout=True, figsize = (19/2.54/3, 11.5/2.54))
        gs = fig.add_gridspec(3,2, width_ratios=widths)

        plot_locs = list(product([0, 1, 2], [0, 1]))[:-1]
        axes = [fig.add_subplot(gs[loc]) for loc in plot_locs]
        cbar_ax = fig.add_subplot(gs[2, 1])
        cbar_axin = inset_axes(cbar_ax,
                            width="30%",  # width = 50% of parent_bbox width
                            height="90%",
                            loc='center left')
        cbar_ax.axis("off")

        for spine in cbar_axin.spines.values():
           spine.set_edgecolor(dark)

        plot_cbars = [False, False, False, False, True]
        cbar_axes = [None, None, None, None, cbar_axin]
        plot_yaxes = [True, False] * 3
        plot_xaxes = [False, False, False, False, True]
        cbar_kwargs = {}
        cbar_kwargs = [None, None, None, None, cbar_kwargs]

        for ds, ax, plot_cbar, cbar_ax, plot_yaxis, plot_xaxis, cbar_kwarg, typ in zip(
                ds_2d, axes, plot_cbars, cbar_axes, plot_yaxes, plot_xaxes, cbar_kwargs,
                range(1,6)
                ):
            ds.plot.contourf(levels=levels, colors=colors, ax=ax,
                             add_colorbar = plot_cbar, cbar_ax=cbar_ax, cbar_kwargs=cbar_kwarg)

            ax.xaxis.label._color = light
            ax.yaxis.label._color = light

            if plot_yaxis == False:
                ax.set_ylabel("")
            else:
                ax.set_ylabel("z [10 m]")
            if plot_xaxis == False:
                ax.set_xlabel("")
                ax.set_xticklabels("")
            else:
                ax.set_xlabel("x [100 km]")

            for spine in ax.spines.values():
                spine.set_edgecolor(dark)

            ax.text(0.1, 0.5, str(typ), horizontalalignment='center',
                 verticalalignment='center', transform=ax.transAxes, color=light,
                 bbox=dict(edgecolor=light, facecolor=dark))

        cbar_axin.set_ylabel("salinity (TDS g/l)", rotation=270, labelpad=10.)
        cbar_axin.xaxis.label._color = light
        cbar_axin.yaxis.label._color = light
        plt.savefig(outf + ".png", facecolor=fig.get_facecolor(), edgecolor='none', dpi=300)
        plt.savefig(outf + ".svg", facecolor=fig.get_facecolor(), edgecolor='none')


"""
第一种方法使用dropwhile()函数
需要注意该函数一旦遇到某个元素不满足测试要求，那么之后的所有元素就会不经过筛选而直接返回
所以下面的例子中文章中间的注释并没有被丢弃
"""
from itertools import dropwhile

with open('../test.txt') as f:
    for line in dropwhile(lambda line: line.startswith('#'), f):
        print(line, end='')

"""
output
Man:what's your name
B:where are you form
>>>>>>>>>>>>>>>>>>>>>>>>>
Man:The question is
Python is powerful
# help
I like Python
It's interesting
"""

"""
第二种方法，如果恰好知道要跳过几个元素，那么可以使用islice()函数
"""
from itertools import islice

x = [1, 2, 3, 'a', 'b', 'c']
for x in islice(x, 3, None): # None表示后面所有元素都包括进来
    print(x, end=' ') # a b c

迭代器和生成器/迭代_跳过前一部分元素_P124.py

迭代_跳过前一部分元素_P124.py


import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
import gc

tokenizer = AutoTokenizer.from_pretrained(
  'kakaobrain/kogpt', revision='KoGPT6B-ryan1.5b-float16',
  bos_token='[BOS]', eos_token='[EOS]', unk_token='[UNK]', pad_token='[PAD]', mask_token='[MASK]'
)

model = AutoModelForCausalLM.from_pretrained(
  'kakaobrain/kogpt', revision='KoGPT6B-ryan1.5b-float16',
  pad_token_id=tokenizer.eos_token_id,
  torch_dtype='auto', low_cpu_mem_usage=True
).to(device='cuda', non_blocking=True)

_ = model.eval()

context = """아래는 친구같은 척척박사 로봇 코밍과 학생의 대화입니다. 역사, 화학과 관련한 질문에 친절하게 답해줍니다.

학생: 동학농민운동이 뭐야?
코밍: 1894년에 일어난 동학농민운동은 반봉건, 반외세의 기치 아래 일어난 농민운동이야. 전봉준, 손화중, 김개남 등이 주도했어.
"""

def generator(prompt, origin_prompt):
    tokens = tokenizer.encode(prompt, return_tensors='pt').to(device='cuda', non_blocking=True)
    gen_tokens = model.generate(
        tokens,
        temperature=0.8,
        max_length=len(tokens[0]) + 4,
        num_beams=4
    )

    generated = tokenizer.batch_decode(gen_tokens)[0]
    if(len(generated[(len(origin_prompt) + 1):]) > 72):
        return generated + "..."

    if generated[(len(origin_prompt) + 1):].find("학생: ") != -1:
        return generated

    return generator(generated, origin_prompt)

logs = []

def getAnswer(question):
    gc.collect()
    torch.cuda.empty_cache()

    with torch.no_grad():
        logs.append("학생: " + question)
        inputs = context + "\n".join(logs[-2:]) + "\n코밍: "
        generated = generator(inputs, inputs)

        say = generated[len(inputs):].split("학생: ")[0].strip().replace("코밍:", "")
        logs.append("코밍: " + say)
        return say


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Wed Feb  8 14:44:35 2023

@author: gabriel
"""

import numpy as np
import matplotlib.pyplot as plt
from hamiltonians import Hamiltonian_A1u_junction_sparse, Hamiltonian_A1u_S, Hamiltonian_A1u_junction_sparse_periodic, Hamiltonian_A1u_S_periodic
from functions import get_components
from phase_functions import phase_soliton_antisoliton_arctan, phase_single_soliton, phase_single_soliton_arctan, phase_soliton_soliton_arctan, phase_soliton_antisoliton, phase_soliton_antisoliton_arctan_A1u_S_around_zero, \
    phase_antisoliton_soliton, phase_soliton_antisoliton_arctan_A1u_S_around_pi, phase_soliton_antisoliton_S_around_zero, phase_soliton_antisoliton_S_around_pi
import scipy

L_x = 200
L_y = 200       #L_y should be odd for single soliton
t = 1
Delta = 0.4*t
mu = -2*t  #-2
t_J = 0.2*t   #t
L = 100      #L_y//2
n = 12 #number of eigenvalues
lambda_J = 10
phi_external = 0
# phi_eq = 0.25*2*np.pi
y = np.arange(1, L_y+1)
y_0 = (L_y-L)//2
y_1 = (L_y+L)//2
y_s = (L_y+1)//2

# Phi = phi_profile(phi_external, y, L_y//2, lambda_J)
# Phi = phase_single_soliton(phi_external, y, y_s)
# Phi = phase_single_soliton_arctan(phi_external, y, y_s, lambda_J)
# Phi = 2*np.pi * np.ones_like(y)
Phi = phase_soliton_antisoliton(phi_external, y, y_0, y_1)
# Phi = phase_soliton_antisoliton_arctan(phi_external, y, y_0, y_1, lambda_J)
# Phi = phase_soliton_antisoliton_arctan_A1u_S_around_zero(phi_external, y, y_0, y_1, lambda_J)
# Phi = phase_antisoliton_soliton(phi_external, y, y_0, y_1)
# Phi = phase_soliton_antisoliton_S_around_pi(phi_external, phi_eq, y, y_0, y_1)
# Phi = phase_soliton_antisoliton_S_around_zero(phi_external, phi_eq, y, y_0, y_1)

params = {"t": t, "mu": mu, "L_x": L_x, "L_y": L_y, "Delta": Delta, "t_J": t_J, "L": L}
H = Hamiltonian_A1u_junction_sparse(t=t, mu=mu, L_x=L_x, L_y=L_y, Delta=Delta, t_J=t_J, Phi=Phi)
# H = Hamiltonian_A1u_junction_sparse_periodic(t=t, mu=mu, L_x=L_x, L_y=L_y, Delta=Delta, t_J=t_J, Phi=Phi)
# H = Hamiltonian_A1u_S(t=t, mu=mu, L_x=L_x, L_y=L_y, Delta=Delta, t_J=t_J, Phi=Phi)
# H = Hamiltonian_A1u_S_periodic(t=t, mu=mu, L_x=L_x, L_y=L_y, Delta=Delta, t_J=t_J, Phi=Phi)
# H = Hamiltonian_A1u_junction_sparse_periodic_in_x(t=t, mu=mu, L_x=L_x, L_y=L_y, Delta=Delta, t_J=t_J, Phi=Phi)
# H = Hamiltonian_A1u_junction(t=t, mu=mu, L_x=L_x, L_y=L_y, Delta=Delta, t_J=t_J, Phi=Phi)

eigenvalues_sparse, eigenvectors_sparse = scipy.sparse.linalg.eigsh(H, k=n, sigma=0)

#%% Probability density
index = np.arange(n)   #which zero mode (less than k)
probability_density = []
zero_state = []
localized_state_upper_left = [] # it is the site (L_x/2-1, (L_y+L)/2)
localized_state_upper_right = [] # it is the site (L_x/2, (L_y+L)/2)
localized_state_bottom_left = [] # it is the site (L_x/2-1, (L_y-L)/2)
localized_state_bottom_right = [] # it is the site (L_x/2, (L_y-L)/2)
localized_state_left = []
localized_state_right = []
probability_density_particle = []

for i in index:
    destruction_up, destruction_down, creation_down, creation_up = get_components(eigenvectors_sparse[:,i], L_x, L_y)
    probability_density.append((np.abs(destruction_up)**2 + np.abs(destruction_down)**2 + np.abs(creation_down)**2 + np.abs(creation_up)**2)/(np.linalg.norm(np.abs(destruction_up)**2 + np.abs(destruction_down)**2 + np.abs(creation_down)**2 + np.abs(creation_up)**2)))
    probability_density_particle.append((np.abs(destruction_up)**2 + np.abs(destruction_down)**2))
    zero_state.append(np.stack((destruction_up, destruction_down, creation_down, creation_up), axis=2)) #positive energy eigenvector splitted in components
    localized_state_upper_left.append(zero_state[i][(L_y+L)//2, L_x//2-1,:])
    localized_state_upper_right.append(zero_state[i][(L_y+L)//2, L_x//2,:])
    localized_state_bottom_left.append(zero_state[i][(L_y-L)//2, L_x//2-1,:])
    localized_state_bottom_right.append(zero_state[i][(L_y-L)//2, L_x//2,:])
    localized_state_left.append(zero_state[i][:, L_x//2-1,:])
    localized_state_right.append(zero_state[i][:, L_x//2,:])

#%% Plotting of probability density

plt.rc("font", family="serif")  # set font family
plt.rc("xtick", labelsize="large")  # reduced tick label size
plt.rc("ytick", labelsize="large")
plt.rc('font', size=18) #controls default text size
plt.rc('axes', titlesize=18) #fontsize of the title
plt.rc('axes', labelsize=18) #fontsize of the x and y labels
plt.rc("text", usetex=True) # for better LaTex (slower)
plt.rcParams['xtick.top'] = True    #ticks on top
plt.rcParams['xtick.labeltop'] = False
plt.rcParams['ytick.right'] = True    #ticks on left
plt.rcParams['ytick.labelright'] = False
plt.rc('legend', fontsize=18) #fontsize of the legend

index = 1
fig, ax = plt.subplots()
image = ax.imshow(probability_density[index], cmap="Blues", origin="lower") #I have made the transpose and changed the origin to have xy axes as usually
plt.colorbar(image)
#ax.set_title(f"{params}")
ax.set_xlabel("x")
ax.set_ylabel("y")
#ax.text(5,25, rf'$index={index}; \Phi={np.round(Phi, 2)}$')
#plt.plot(probability_density[10,:,0])
ax.set_title("Probability density")
plt.tight_layout()
# probability_density_right = probability_density[index][:, L_x//2]/np.linalg.norm(probability_density[0][:, L_x//2])  #The y-axis is inverted
probability_density_right = probability_density[index][:, L_x-1]/np.linalg.norm(probability_density[index][:, L_x-1])  #The y-axis is inverted

fig, ax = plt.subplots()
ax.plot(y, probability_density_right, "o")
#ax.plot(np.arange(1, L_y+1), probability_density[index][:, L_x//2-1])
ax.set_xlabel(r"$\ell$")
ax.set_ylabel("Probability density at the junction")
ax.text(5,25, rf'$index={index}$')
ax.set_xticks([1,50,100,150,200])
# ax.set_title("Probability density at the junction")
# plt.tight_layout()

fig, ax = plt.subplots()
plt.title("Probability density particle")
image = ax.imshow(probability_density_particle[index], cmap="Blues", origin="lower") #I have made the transpose and changed the origin to have xy axes as usually
plt.colorbar(image)
#ax.set_title(f"{params}")
ax.set_xlabel("x")
ax.set_ylabel("y")

fig, ax = plt.subplots()
plt.title("Soliton profile")
ax.plot(y, Phi)
ax.set_xlabel("x")
ax.set_ylabel("y")

#%% Spin determination
from functions import mean_spin_xy, get_components

# zero_modes = eigenvectors_sparse
# creation_up, creation_down, destruction_down, destruction_up = get_components(zero_modes[:,index], L_x, L_y)
# zero_state = np.stack((destruction_up, destruction_down, creation_down, creation_up), axis=2) #positive energy eigenvector splitted in components
#corner_state = zero_state[L, L_y//2, :].reshape(4,1)  #positive energy point state localized at the junction
# corner_state_normalized = corner_state/np.linalg.norm(corner_state[:2]) #normalization with only particle part
# spin_mean_value = mean_spin(corner_state_normalized)

spin = []
for i in range(n):
    spin.append(mean_spin_xy(zero_state[i]))

#%%
# fig, ax = plt.subplots()
# image = ax.imshow(spin[:,:,2].T, cmap="Blues", origin="lower") #I have made the transpose and changed the origin to have xy axes as usually
# plt.colorbar(image)
#image.set_clim(np.min(spin[:,:,1].T), np.max(spin[:,:,1].T))

# Meshgrid
x_mesh, y_mesh = np.meshgrid(np.linspace(0, L_x-1, L_x),
                    #np.linspace(L_y-1, 0, L_y))
                    np.linspace(0, L_y-1, L_y))

# Directional vectors
u = spin[index][:, :, 0]   #x component
v = spin[index][:, :, 1]   #y component

# Plotting Vector Field with QUIVER
fig,ax = plt.subplots()
ax.quiver(x_mesh, y_mesh, u, v, color='r', angles='uv')
ax.set_title('Spin Field in the plane')

#%% Spin in z
fig, ax = plt.subplots()
image = ax.imshow(spin[index][:,:,2], cmap="Blues", origin="lower") #I have made the transpose and changed the origin to have xy axes as usually
plt.colorbar(image)
ax.set_title('Spin Field in the z direction')
plt.text(0,0, f"index={index}")

# fig, ax = plt.subplots()
# ax.plot(spin[:, L_x//2,2])
# total_spin = np.sum(spin[:, L_x//2, 2])
# plt.text(0,0.25, f"Total spin={total_spin}, index={index}")
# plt.text(0,-0.25, f"Total spin={total_spin}, index={index}")
# ax.set_title('Spin Field in the z direction')
# plt.text(0,0, f"index={index}")

#%% Energy spectrum

fig, ax = plt.subplots()
plt.plot(eigenvalues_sparse, "o")
ax.set_xlabel("Label of eingevalue")
ax.set_ylabel("Energy")
plt.tight_layout()

#%% Phi spectrum
"""
from functions import phi_spectrum_sparse_single_step

Phi_values = np.linspace(0, 2*np.pi, 10)
phi_energy = phi_spectrum_sparse_single_step(Hamiltonian_A1u_single_step_sparse, Phi_values, t, mu, L_x, L_y, Delta, t_J)

fig, ax = plt.subplots()
ax.plot(Phi_values, phi_energy)
ax.set_xlabel(r"$\phi$")
ax.set_ylabel("E")

"""

#%% Spinors to txt
from functions import mean_spin

index = range(0,4)

with open("spinors.txt", "w+") as f:
    data = f.read()
    f.write(f"{params}\n")
    f.write(f"energies={eigenvalues_sparse}\n\n")
    for i in index:
        f.write(f"{i}th-localized state at the upper soliton\n\n")
        for j in range(4):
            f.write(f"{str(localized_state_upper_left[i].round(4)[j]):30}"+"%    "+
                    f"{str(localized_state_upper_right[i].round(4)[j])}"+"\n")
        f.write("\n")
        f.write("%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%\n")
    f.write("%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%\n")
    f.write("%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%\n")
    for i in index:
        f.write(f"{i}th-localized state at the bottom soliton\n\n")
        for j in range(4):
            f.write(f"{str(localized_state_bottom_left[i].round(4)[j]):30}"+"%    "+
                    f"{str(localized_state_bottom_right[i].round(4)[j])}"+"\n")
        f.write("\n")
        f.write("%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%\n")


import jsonpickle

from cloudshell.cp.docker.common.deploy_data_holder import DeployDataHolder
from cloudshell.cp.docker.models.deploy_docker_container_resource_model import DeployDockerContainerResourceModel
from cloudshell.cp.docker.models.docker_cloud_provider_resource_model import DockerCloudProviderResourceModel

class DockerModelsParser(object):
    @staticmethod
    def convert_app_resource_to_deployed_app(resource):
        json_str = jsonpickle.decode(resource.app_context.deployed_app_json)
        data_holder = DeployDataHolder(json_str)

        return data_holder

    @staticmethod
    def convert_to_docker_resource_model(resource):
        resource_context = resource.attributes

        docker_resource_model = DockerCloudProviderResourceModel()

        return docker_resource_model

    @staticmethod
    def convert_to_deployment_resource_model(deployment_request):
        data = jsonpickle.decode(deployment_request)

        data_holder = DeployDataHolder(data)

        deployment_resource_model = DeployDockerContainerResourceModel()
        deployment_resource_model.app_name = data_holder.app_name
        deployment_resource_model.cloud_provider = data_holder.docker_params.cloud_provider
        deployment_resource_model.docker_image = data_holder.docker_params.docker_image
        deployment_resource_model.docker_image_command = data_holder.docker_params.docker_image_command
        deployment_resource_model.docker_image_tag = data_holder.docker_params.docker_image_tag

        return deployment_resource_model

    @staticmethod
    def convert_to_bool(string):
        """
        Converts string to bool
        :param string: String
        :str string: str
        :return: True or False
        """
        if isinstance(string, bool):
            return string
        return string in ['true', 'True', '1']

    @staticmethod
    def get_public_ip_from_connected_resource_details(resource_context):
        public_ip_on_resource = ""
        public_ip = 'Public IP'
        if resource_context.remote_endpoints is not None and public_ip in resource_context.remote_endpoints[
            0].attributes:
            public_ip_on_resource = resource_context.remote_endpoints[0].attributes[public_ip]
        return public_ip_on_resource

    @staticmethod
    def get_private_ip_from_connected_resource_details(resource_context):
        private_ip_on_resource = ""
        if resource_context.remote_endpoints is not None:
            private_ip_on_resource = resource_context.remote_endpoints[0].address
        return private_ip_on_resource

    @staticmethod
    def try_get_deployed_connected_resource_instance_id(resource_context):
        try:
            deployed_instance_id = str(
                jsonpickle.decode(resource_context.remote_endpoints[0].app_context.deployed_app_json)['vmdetails'][
                    'uid'])
        except Exception as e:
            raise ValueError('Could not find an ID of the Docker Deployed instance' + e.message)
        return deployed_instance_id

    @staticmethod
    def get_connectd_resource_fullname(resource_context):
        if resource_context.remote_endpoints[0]:
            return resource_context.remote_endpoints[0].fullname
        else:
            raise ValueError('Could not find resource fullname on the deployed app.')


from django.urls import path
from .views import (
    PostListView,
    UserPostListView
)

from . import views

urlpatterns = [
    path('',views.index, name = 'index'),
    path('user/<str:username>',UserPostListView.as_view(), name = 'speech-posts'),
    path('speech/', PostListView.as_view(), name = 'speech'),
    path('about/', views.about, name = 'EndOf-about'),
    path('about/', views.about, name = 'EndOf-about'),
    path('speaker/',views.Speakview, name = 'speaker_detail'),
]


# Create your views here.
from django.shortcuts import render
from rest_framework.response import Response
from rest_framework.views import APIView
from rest_framework import status
from django.contrib.auth import authenticate, login, logout
from .models import User
from .serializers import UserSignUpSerializer, UserLoginSerializer

class SignUpApiView(APIView):
    def post(self, request, *args, **kwargs):
        data = {}
        serializer = UserSignUpSerializer(data=request.data)
        if serializer.is_valid(raise_exception=True):
            user = serializer.save()
            data['username'] = user.username
            data['email'] = user.email
            return Response(data, status=status.HTTP_200_OK)
        else:
            data['error'] = "error"
            return Response(data, status=status.HTTP_404_NOT_FOUND)

class LoginApiView(APIView):

    def post(self, request, *args, **kwargs):
        data = {}
        serializer = UserLoginSerializer(data=request.data)
        import pdb;pdb.set_trace()
        if serializer.is_valid():
            authenticated_user = serializer.validated_data.get('authenticated_user')
            login(request, authenticated_user)
            data['id'] = authenticated_user.id
            data['username'] = authenticated_user.username
            return Response(data, status=status.HTTP_200_OK)
        else:
            data['error'] = serializer.error
            return Response(data, status=status.HTTP_404_NOT_FOUND)

class LogoutApiView(APIView):

    queryset = []
    def get(self, request):
        data = {
            'msg': "logged out successfully."
        }
        logout(request)
        return Response(data, status=status.HTTP_200_OK)


import numpy as np
import regions
from astropy import units as u
from astropy.modeling import models, fitting
from astropy import wcs
from astropy import coordinates
from astropy.io import fits
from astropy import convolution
from spectral_cube import SpectralCube
import pvextractor
from gaussfit_catalog.core import gaussfit_image
import paths
import os
import pylab as pl
import shapely.geometry as geom
from show_pv import show_keplercurves, show_pv
import imp
import edge_on_ring_velocity_model
from astropy.utils.console import ProgressBar
from constants import vcen as assumed_vcen, origin, d_orion
import tracemalloc
from astropy import log

imp.reload(edge_on_ring_velocity_model)
from edge_on_ring_velocity_model import thindiskcurve, thindiskcurve_fitter, trace_plotter

first_snapshot = tracemalloc.take_snapshot()

if 'cached_gaussfit_results' not in locals():
    cached_gaussfit_results = {}

fitresult_list = []

for linename,(vmin,vmax),limits,(cenx, ceny) in (
    #('Unknown_4', (-15, 27), (-0.1, 0.1, -0.12, 0.12), (65.1, 60.5)),
    #('H2Ov2=1_5(5,0)-6(4,3)', (-28, 38), (-0.2, 0.2, -0.2, 0.2), (68.8, 65.5)),
    #('Unknown_1', (-15, 27), (-0.1, 0.1, -0.12, 0.12), (65.1, 60.5)),
    #('Unknown_2', (-15, 27), (-0.1, 0.1, -0.12, 0.12), (65.1, 60.5)),
    ('SiOv=1_5-4', (-30, 45), (-0.2, 0.2, -0.2, 0.2), (65.1, 60.5)),
   ):

    regs = regions.Regions.read(paths.rpath('velo_centroid_guesses_{linename}.reg').format(linename=linename))

    guesses = {}

    for reg in regs:
        vel = float(reg.meta['text'].strip("{}")
                    if 'text' in reg.meta
                    else reg.meta['label']
                   )
        if vel in guesses:
            guesses[vel].append(reg)
        else:
            guesses[vel] = [reg]

    velocities = np.array(sorted(guesses.keys()))*u.km/u.s

    cubefn = paths.dpath('cubes/OrionSourceI_{linename}_robust0.5maskedclarkclean10000_medsub_K.fits'
                         .format(linename=linename))
    basename = os.path.splitext(os.path.basename(cubefn))[0]
    cube = SpectralCube.read(cubefn)

    # hard-coded icrs radecsys assumed
    assert cube.header['RADESYS'] == 'ICRS'

    vdiff = np.abs(np.diff(cube.spectral_axis).mean())
    rms = cube.std()
    weights = np.ones(cube.shape[1:]) * rms.value**2

    try:
        beam = cube.beam
    except AttributeError:
        beam = cube.average_beams(1)
    pixscale = wcs.utils.proj_plane_pixel_area(cube.wcs.celestial)**0.5 * u.deg
    bmmaj_px = (beam.major / pixscale).decompose().value
    bmmin_px = (beam.minor / pixscale).decompose().value
    max_radius_in_beams = 1.25
    max_offset_in_beams = 1.5

    STDDEV_TO_FWHM = np.sqrt(8*np.log(2))

    results = {}

    if linename in cached_gaussfit_results:
        print("Loading {0} from cache".format(linename))
        results = cached_gaussfit_results[linename]
    else:
        print("Fitting {0}, which is not in cache".format(linename))
        pb = ProgressBar(cube.shape[0])

        for vel,vslice in (zip(cube.spectral_axis, cube)):
            closest = np.argmin(np.abs(vel-velocities))
            if np.abs(velocities[closest] - vel) > vdiff:
                #print("Skipping velocity {0}, closest is {1} -> {2}".format(vel, closest,
                #                                                            velocities[closest]))
                pb.update()
                continue

            thisvel = velocities[closest].value
            guess_regs = guesses[thisvel]

            ampguess = vslice.max().value

            model_list = []
            for reg in guess_regs:

                p_init = models.Gaussian2D(amplitude=ampguess,
                                           x_mean=reg.center.x,
                                           y_mean=reg.center.y,
                                           x_stddev=bmmaj_px/STDDEV_TO_FWHM*0.75,
                                           y_stddev=bmmin_px/STDDEV_TO_FWHM*0.75,
                                           theta=beam.pa.to(u.deg).value,
                                           bounds={'x_stddev':(bmmin_px/STDDEV_TO_FWHM*0.5,
                                                               bmmaj_px*max_radius_in_beams/STDDEV_TO_FWHM),
                                                   'y_stddev':(bmmin_px/STDDEV_TO_FWHM*0.5,
                                                               bmmaj_px*max_radius_in_beams/STDDEV_TO_FWHM),
                                                   'x_mean':(reg.center.x-max_offset_in_beams*bmmaj_px/STDDEV_TO_FWHM,
                                                             reg.center.x+max_offset_in_beams*bmmaj_px/STDDEV_TO_FWHM),
                                                   'y_mean':(reg.center.y-max_offset_in_beams*bmmaj_px/STDDEV_TO_FWHM,
                                                             reg.center.y+max_offset_in_beams*bmmaj_px/STDDEV_TO_FWHM),
                                                   'amplitude':(ampguess*0.1, ampguess*2.1)
                                                  }
                                          )
                assert not p_init._has_units
                model_list.append(p_init)

            composite_model = model_list[0]
            if len(model_list) > 1:
                for mod in model_list[1:]:
                    composite_model += mod

            fit_result = gaussfit_image(vslice.value, composite_model,
                                        weights=weights, plot=True)
            results[thisvel] = fit_result

            if not os.path.exists(paths.fpath('velcentroid/diagnostics/{0}/'.format(basename))):
                os.mkdir(paths.fpath('velcentroid/diagnostics/{0}/').format(basename))
            pl.savefig(paths.fpath('velcentroid/diagnostics/{0}/{1}.png')
                       .format(basename, thisvel))

            pb.update()

        cached_gaussfit_results[linename] = results

    diskend_regs = regions.Regions.read(paths.rpath('diskends.reg'))
    diskends = coordinates.SkyCoord([reg.center for reg in diskend_regs]).icrs

    #center = coordinates.SkyCoord(diskends.ra.mean(), diskends.dec.mean(),
    #                              frame=diskends.frame)
    center_cont = regions.Regions.read(paths.rpath('sourceI_center.reg'))[0].center.transform_to(coordinates.ICRS)
    center_U = regions.Regions.read(paths.rpath('sourceI_center.reg'))[1].center.transform_to(coordinates.ICRS)
    center_h2o = regions.Regions.read(paths.rpath('sourceI_center.reg'))[2].center.transform_to(coordinates.ICRS)

    center = center_cont

    ref_cen_x, ref_cen_y = cube.wcs.celestial.wcs_world2pix(center.ra.deg,
                                                            center.dec.deg, 0)
    # had to switch to pixel centering because coordinate transformations
    # aren't accurate enough (and/or, don't match ds9)
    ref_cen_x = cenx
    ref_cen_y = ceny

    #assert np.abs(ref_cen_x - 66) < 2
    #assert np.abs(ref_cen_y - 63) < 2

    def offset_to_point(xx, yy):
        line = geom.LineString(zip(diskends.ra.deg, diskends.dec.deg))
        point = geom.Point(xx, yy)
        return line.project(point)

    ref_offset = offset_to_point(center.ra.deg, center.dec.deg)

    # Add the data as a background image
    pvfn = 'pv/sourceI_{0}_B6_robust0.5_diskpv_0.01.fits'.format(linename)
    fh = fits.open(paths.dpath(pvfn))
    data = fh[0].data
    header = fh[0].header

    ww = wcs.WCS(fh[0].header)

    ww.wcs.cdelt[1] /= 1000.0
    ww.wcs.crval[1] /= 1000.0
    ww.wcs.cunit[1] = ' '#u.km/u.s
    ww.wcs.cdelt[0] *= 3600
    ww.wcs.cunit[0] = u.arcsec
    ww.wcs.crval[0] = -origin.to(u.arcsec).value

    # #!@#$!@#$@!%@#${^(@#$)%#$(
    ww.wcs.set()

    snap3 = tracemalloc.take_snapshot()
    diff = snap3.compare_to(first_snapshot, 'lineno')
    mem_used = sum([dd.size_diff for dd in diff])*u.B
    log.info(f"Used {mem_used.to(u.GB)} for line {linename} prior to plotting")

    fig1 = pl.figure(1)
    pl.clf()
    ax1 = fig1.add_subplot(1,2,1)
    pvwcs = pvextractor.utils.wcs_slicing.slice_wcs(cube.wcs, pixscale)
    pvwcs.wcs.cdelt[0] *= 3600
    pvwcs.wcs.crval[0] = -origin.to(u.arcsec).value
    pvwcs.wcs.cunit[0] = u.arcsec
    pvwcs.wcs.crval[1] /= 1e3
    # not clear why this is needed?
    pvwcs.wcs.crval[1] += assumed_vcen.to(u.km/u.s).value
    pvwcs.wcs.cdelt[1] /= 1e3
    pvwcs.wcs.cunit[1] = ' ' # force wcs.set to NOT revert to m/s
    print('pvwcs: cdelt={0} crval={1} crpix={2}'.format(pvwcs.wcs.cdelt,
                                                        pvwcs.wcs.crval,
                                                        pvwcs.wcs.crpix))
    ax2 = fig1.add_subplot(1,2,2, projection=ww)

    trans = ax2.get_transform('world')

    diskends_x, diskends_y = cube.wcs.celestial.wcs_world2pix(diskends.ra.deg,
                                                              diskends.dec.deg,
                                                              0)
    ax1.plot((diskends_x-ref_cen_x)*pixscale.to(u.arcsec),
             (diskends_y-ref_cen_y)*pixscale.to(u.arcsec), 'k-',
             linewidth=3, alpha=0.2, zorder=-100)
    ax1.plot(0, 0, 'ko', markersize=5, alpha=0.5, zorder=-50)

    loc = pl.matplotlib.ticker.MultipleLocator(base=0.04)
    ax1.xaxis.set_major_locator(loc)

    cmap = pl.cm.Spectral_r
    cmap = pl.cm.Spectral
    norm = pl.matplotlib.colors.Normalize(vmin=vmin, vmax=vmax)

    offset_fits_arcsec = []
    vels = np.sort(np.array([v for v in results]))

    gfit_points = []

    for vel in vels:
        fit_result = results[vel]
        fitted = fit_result[0]
        n_submodels = fitted.n_submodels()
        color = cmap(norm(vel))
        if n_submodels > 1:
            offsets_ = []
            for ii in range(n_submodels):
                xcen, ycen = (getattr(fitted, 'x_mean_{0}'.format(ii)),
                              getattr(fitted, 'y_mean_{0}'.format(ii)))
                ax1.plot((xcen-ref_cen_x)*pixscale.to(u.arcsec),
                         (ycen-ref_cen_y)*pixscale.to(u.arcsec),
                         color=color,
                         marker='o')
                ra,dec = cube.wcs.celestial.wcs_pix2world(xcen, ycen, 0)
                offset = offset_to_point(ra, dec)
                gfit_points.extend(ax2.plot((offset-ref_offset)*3600, vel,
                                            color=color, marker='s',
                                            markersize=5,
                                            markeredgecolor='k',
                                            transform=trans))
                offsets_.append((offset-ref_offset)*3600)
            offset_fits_arcsec.append(np.mean(offsets_))
        else:
            xcen, ycen = fitted.x_mean, fitted.y_mean
            ax1.plot((xcen-ref_cen_x)*pixscale.to(u.arcsec),
                     (ycen-ref_cen_y)*pixscale.to(u.arcsec), color=color, marker='o')
            ra,dec = cube.wcs.celestial.wcs_pix2world(xcen, ycen, 0)
            offset = offset_to_point(ra, dec)
            gfit_points.extend(ax2.plot((offset-ref_offset)*3600, vel,
                                        color=color, marker='s',
                                        markersize=5,
                                        markeredgecolor='k',
                                        transform=trans)
                              )
            offset_fits_arcsec.append((offset-ref_offset)*3600)

    show_keplercurves(ax2, 0*u.deg, 150, assumed_vcen, yaxis_unit=u.km/u.s, radii={},
                      masses=[15, ], linestyles=[':',], colors=['r'],
                      trans=ax2.get_transform('world'),
                     )

    # xx_thindisk, yy_thindisk = thindiskcurve(mass=20*u.M_sun, rmin=30*u.au, rmax=80*u.au)
    # ax2.plot((xx_thindisk / d_orion).to(u.arcsec, u.dimensionless_angles()),
    #          yy_thindisk + assumed_vcen,
    #          'k:',
    #          transform=trans)
    xx_thindisk, yy_thindisk = thindiskcurve(mass=15.5*u.M_sun, rmin=17*u.au,
                                             rmax=66*u.au)
    thindiskline = ax2.plot((xx_thindisk / d_orion).to(u.arcsec,
                                                       u.dimensionless_angles()),
                            yy_thindisk + assumed_vcen, 'k-', transform=trans)

    xlim, ylim = pvwcs.wcs_world2pix([-0.2,0.2], [vmin-3, vmax+3], 0)
    xlim, ylim = ww.wcs_world2pix([-0.2,0.2], [vmin-3, vmax+3], 0)
    ax2.set_ylim(*ylim)
    ax2.set_xlim(*xlim)
    ax2.yaxis.tick_right() # mpl version; incompatible with wcsaxes
    ax2.coords[1].set_ticklabel_position('r')
    ax2.coords[1].set_axislabel_position('r')
    ax2.set_ylabel("$V_{LSR}$ [km s$^{-1}$]")
    ax2.set_xlabel("Offset Position (arcsec)")
    ax1.set_xlabel("Offset RA (arcsec)")
    ax1.set_ylabel("Offset Dec (arcsec)")
    ax1.axis(limits)

    for tick in ax1.get_xticklabels():
        tick.set_rotation(90)
    for tick in ax2.get_xticklabels():
        tick.set_rotation(90)

    pl.savefig(paths.fpath('velcentroid/{linename}_pp_pv_plots.pdf'.format(linename=linename)),
               bbox_inches='tight')

    im = ax2.imshow(data, cmap='gray_r', interpolation='none', vmax=0.02,
                    vmin=-0.0001, origin='lower',
                    #transform=ax2.get_transform('world'),
                   )
    ax2.set_aspect('auto')
    ax2.set_ylim(*ylim)
    ax2.set_xlim(*xlim)

    pl.savefig(paths.fpath('velcentroid/{linename}_pp_pv_plots_withbackground.pdf'.format(linename=linename)),
               bbox_inches='tight')

    #fig3,ax3 = show_pv(data, ww,
    #                   origin, vrange=[-40,55], vcen=assumed_vcen,
    #                   imvmin=-0.001, imvmax=0.020)

    #trans3 = ax3.get_transform('world')
    #overlay_thindiskline = ax3.plot((xx_thindisk / d_orion).to(u.arcsec,
    #                                                           u.dimensionless_angles()),
    #                                 (yy_thindisk + assumed_vcen).to(u.m/u.s), 'k-', transform=trans3)

    # previously had an
    #if 'Unknown' in linename:
    # statement here because these are the only ones that represent disks,
    # but since we're using the *average* offset position, maybe this is OK?

    offsets_au = (u.Quantity(offset_fits_arcsec, u.arcsec) *
                  d_orion).to(u.au, u.dimensionless_angles())
    print("Beginning thin disk fitting for {0}".format(linename))
    if False:
        fitresult, ci, trace, minimizer = thindiskcurve_fitter(
            xsep=np.array(offsets_au), velo=vels, mguess=15.5*u.M_sun, rinner=17,
            router=66, conf_int=True, npix=500)
        fitresult_list.append(fitresult)
        trace_plotter(ci, trace, fitresult, minimizer)
    else:
        fitresult = thindiskcurve_fitter(xsep=np.array(offsets_au), velo=vels,
                                         mguess=15.5*u.M_sun, rinner=17,
                                         router=66, conf_int=False, npix=500)
        fitresult_list.append(fitresult)

    for line in thindiskline:
        line.set_visible(False)

    xx_thindisk, yy_thindisk = thindiskcurve(mass=fitresult.params['mass']*u.M_sun,
                                             rmin=fitresult.params['rinner']*u.au,
                                             rmax=fitresult.params['router']*u.au,
                                            )

    lines = ax2.plot((xx_thindisk / d_orion).to(u.arcsec, u.dimensionless_angles()),
                     yy_thindisk + assumed_vcen,
                     'k-',
                     transform=trans,
                     label=('$M={0:0.1f}$\n$R_{{in}}={1:d}$\n$R_{{out}}={2:d}$'
                            .format(fitresult.params['mass'].value,
                                    int(fitresult.params['rinner'].value),
                                    int(fitresult.params['router'].value),
                                   )
                           )
                    )

    xmin,xmax = ax2.get_xlim()
    ax2.hlines(fitresult.params['vcen'].value, xmin, xmax, linestyle='--', color='k',
               alpha=0.5, zorder=500, transform=trans)

    leg = pl.legend(loc='upper left', fontsize=12, handlelength=1.0)

    im.set_visible(False)

    pl.savefig(paths.fpath('velcentroid/{linename}_pp_pv_plots_fittedmodel.pdf'.format(linename=linename)),
               bbox_inches='tight')

    # show the average fitted positions
    avgposns = ax2.plot(offset_fits_arcsec, vels, marker='o', markersize=5,
                        linestyle='none', markerfacecolor='k', alpha=0.5,
                        markeredgecolor='w', transform=trans)

    pl.savefig(paths.fpath('velcentroid/{linename}_pp_pv_plots_fittedmodel_withavgs.pdf'.format(linename=linename)),
               bbox_inches='tight')

    im.set_visible(True)
    pl.savefig(paths.fpath('velcentroid/{linename}_pp_pv_plots_fittedmodel_withavgs_withbackground.pdf'.format(linename=linename)),
               bbox_inches='tight')
    im.set_visible(False)

    tracemalloc.start()

    snap1 = tracemalloc.take_snapshot()

    for mass in (15.5, 19, 5, 10, 15, 20, ):
        for line in lines:
            line.remove()
            #line.set_visible(False)
        leg.remove()
        del leg, lines

        fitresult = thindiskcurve_fitter(xsep=np.array(offsets_au),
                                         velo=vels,
                                         mguess=mass*u.M_sun,
                                         rinner=17,
                                         router=66,
                                         fixedmass=True,
                                        )
        assert fitresult.params['mass'].value == mass

        snap2 = tracemalloc.take_snapshot()
        diff = snap2.compare_to(snap1, 'lineno')
        mem_used = sum([dd.size_diff for dd in diff])*u.B
        log.info(f"Used {mem_used.to(u.GB)} with mass={mass} after fitresult")

        xx_thindisk, yy_thindisk = thindiskcurve(mass=fitresult.params['mass']*u.M_sun,
                                                 rmin=fitresult.params['rinner']*u.au,
                                                 rmax=fitresult.params['router']*u.au,
                                                )

        lines = ax2.plot((xx_thindisk / d_orion).to(u.arcsec, u.dimensionless_angles()),
                         yy_thindisk + assumed_vcen,
                         'k-',
                         transform=trans,
                         label=('$M={0:0.1f}$\n$R_{{in}}={1:d}$\n$R_{{out}}={2:d}$'
                                .format(fitresult.params['mass'].value,
                                        int(fitresult.params['rinner'].value),
                                        int(fitresult.params['router'].value),
                                       )
                               )
                        )
        leg = pl.legend(lines,
                        [line.get_label() for line in lines],
                        loc='upper left',
                        fontsize=12, handlelength=1.0)
        pl.tight_layout()
        pl.savefig(paths.fpath('velcentroid/{linename}_pp_pv_plots_fittedmodel_{mass}msun_withavgs.pdf'
                               .format(mass=mass, linename=linename)),
                   bbox_inches='tight'
                  )

        del fitresult, xx_thindisk, yy_thindisk

        snap2 = tracemalloc.take_snapshot()
        diff = snap2.compare_to(snap1, 'lineno')
        mem_used = sum([dd.size_diff for dd in diff])*u.B
        log.info(f"Used {mem_used.to(u.GB)} with mass={mass}")

    im.set_visible(True)
    for line in lines:
        line.set_visible(False)
    leg.remove()

    xx,vv,pvd = thindiskcurve(mass=15*u.M_sun,
                              rmin=25*u.au,
                              rmax=65*u.au,
                              pvd=True,
                             )
    xx_as = (xx/d_orion).to(u.arcsec, u.dimensionless_angles())

    conv_beam = convolution.Gaussian2DKernel(0.04*u.arcsec/np.diff(xx_as).mean(),
                                             1.5*u.km/u.s/np.diff(vv).mean(),
                                            )
    sm_pvd = convolution.convolve_fft(pvd, conv_beam)

    ax1.cla()
    #ax1.remove()
    #ax1 = fig1.add_subplot(1,2,1, projection=ww)
    ax1.imshow(sm_pvd, cmap='gray_r', extent=[xx_as.value.min(),
                                              xx_as.value.max(),
                                              (vv+assumed_vcen).value.min(),
                                              (vv+assumed_vcen).value.max()],
               interpolation='none', origin='lower',
               vmin=0,
               vmax=sm_pvd[np.abs(vv)<20*u.km/u.s].max(),
               label=('$M={0:0.1f}$\n$R_{{in}}={1:d}$\n$R_{{out}}={2:d}$'
                      .format(15, 25, 65))
              )

    keplines = show_keplercurves(ax1, 0*u.deg, 150, assumed_vcen,
                                 yaxis_unit=u.km/u.s, radii={}, masses=[15, ],
                                 linestyles=[':',], colors=['r'],
                                 trans=ax1.transData,
                     )

    xlim_as, ylim_kms = [-0.2,0.2], [vmin-3, vmax+3]
    ax1.set_xlim(*xlim_as)
    ax1.set_ylim(*ylim_kms)

    ax1.set_aspect('auto')
    #ax1.legend(('$M={0:0.1f}$\n$R_{{in}}={1:d}$\n$R_{{out}}={2:d}$'
    #                  .format(15, 25, 65)),
    #           [im],
    #           loc='best')

    loc = pl.matplotlib.ticker.MultipleLocator(base=0.1)
    ax1.xaxis.set_major_locator(loc)

    ax1.set_ylabel("$V_{LSR}$ [km s$^{-1}$]")
    ax1.set_xlabel("Offset Position (arcsec)")

    pl.savefig(paths.fpath('velcentroid/{linename}_pp_pv_plots_fittedmodel_{mass}msun_withavgs_comparepv.pdf'
                           .format(mass=15, linename=linename)),
               bbox_inches='tight')

    for feature in gfit_points + avgposns:
        feature.set_visible(False)

    pl.savefig(paths.fpath('velcentroid/{linename}_pp_pv_plots_fittedmodel_{mass}msun_comparepv.pdf'
                           .format(mass=15, linename=linename)),
               bbox_inches='tight')
    pl.savefig(paths.fpath('velcentroid/{linename}_pp_pv_plots_fittedmodel_{mass}msun_comparepv.png'
                           .format(mass=15, linename=linename)),
               bbox_inches='tight')

    fig1.clf()

    snap3 = tracemalloc.take_snapshot()
    diff = snap3.compare_to(first_snapshot, 'lineno')
    mem_used = sum([dd.size_diff for dd in diff])*u.B
    log.info(f"Used {mem_used.to(u.GB)} for line {linename}")


import os
import argparse
import models
import utils
import torch
import pandas as pd
import numpy as np
import multiprocessing
from torchvision import datasets
from torchvision import transforms
from tqdm import tqdm

im_transform = transforms.Compose([
    transforms.Resize(64),
    transforms.ToTensor(),
])

parser = argparse.ArgumentParser()
parser.add_argument('--data-type', default='test', type=str,
                    choices=['train', 'test', 'val'])
parser.add_argument('--beta', default=1, type=float)
parser.add_argument('--latent-size', default=16, type=int)
parser.add_argument('--batch-size', default=128, type=int)
parser.add_argument('--learning-rate', default=0.01, type=float)
parser.add_argument('--epochs', default=300, type=float)
parser.add_argument('--image-path',
                    default='../../data/letters/train/')
parser.add_argument('--model-path',
                    default='../../trained_models/checkpoints/')
parser.add_argument('--output-path',
                    default='../../output/')
parser.add_argument('--use-cuda', default=False, action='store_true')
args = parser.parse_args()

hyperparams = utils.dict2string(args.__dict__,
                                ['beta', 'latent_size', 'batch_size',
                                 'learning_rate', 'epochs'])

use_cuda = args.use_cuda and torch.cuda.is_available()
device = torch.device("cuda" if use_cuda else "cpu")
print('Using device', device)
print('num cpus:', multiprocessing.cpu_count())

im_transform = transforms.Compose([
    transforms.Resize(64),
    transforms.ToTensor(),
])

def get_imgs(target_data_type, list_imgs):
    data_types = [os.path.basename(os.path.dirname(os.path.dirname(l[0]))) \
                  for l in list_imgs]

    IXs = []
    for IX, data_type in enumerate(data_types):
        if data_type == target_data_type:
            IXs.append(IX)
    return IXs

# LOAD MODEL
model = models.BetaVAE(latent_size=args.latent_size,
                       beta = args.beta).to(device)
start_epoch = model.load_last_model(args.model_path, hyperparams) + 1

# LOAD DATA
data_dir = '../../data/letters/'
kwargs = {'num_workers': multiprocessing.cpu_count(),
          'pin_memory': True} if use_cuda else {}
dataset = datasets.ImageFolder(os.path.join(data_dir, args.data_type),
                               transform = im_transform)
data_loader = torch.utils.data.DataLoader(dataset, batch_size=1,
                                          shuffle=False, **kwargs)

# EXTRACT ACTIVATIONS
df = []
for i_data, (data, fn_img) in enumerate(zip(data_loader, dataset.imgs)):
    if i_data % 1000 == 0:
        print(f'image {i_data}/{len(data_loader)}')
    _, mu, logvar = model(data[0])
    _, letter, _, size, _, xshift, _, yshift, _, font, _, upper, _ = \
        os.path.basename(fn_img[0]).split('_')
    # ADPPEND TO DATAFRAME
    df.append([letter, size, xshift, yshift, font, upper, mu, logvar])
df = pd.DataFrame(df, columns=['letter', 'size', 'xshift', 'yshift', 'font', 'upper', 'mu', 'logvar'])

# SAVE
output_path = os.path.join(args.output_path, hyperparams[:-1])
os.makedirs(output_path, exist_ok=True)
fn_df = f'activations_{args.data_type}_{hyperparams[:-1]}.csv'
df.to_csv(os.path.join(args.output_path, fn_df))


# Zinan Xu 20871816 and Shengquan Ni 46564157.  ICS 31 Lab sec 7.  Lab asst 8.
from collections import namedtuple

def P(argu=''):
    if not isinstance(argu,list):
        if not isinstance(argu,str):
            argu=str(argu)
        if "\n" in argu:
            argu=argu.replace("\n","\n    ")
        print("    "+argu)
    else:
        P_list(argu) if len(argu)!=0 else print("    empty list")

def P_list(list):
    for i in list:
        P(i)

P('\n----------part(c)----------\n')
P('\n----------part(c.1)----------\n')

Dish = namedtuple('Dish', 'name price calories')
def read_menu_with_count(filepath)->list:
    result=[]
    file=open(filepath)
    count=int(file.readline().replace('\n',''))
    if count:
        for i in range(count):
            fields=file.readline().replace('\n','').split('\t')
            result.append(Dish(fields[0],float(fields[1][1:]),float(fields[2])))
    return result
P('menu1:')
P(read_menu_with_count('menu1.txt'))
P('menu2:')
P(read_menu_with_count('menu2.txt'))

P('\n----------part(c.2)----------\n')

def read_menu(filepath)->list:
    result=[]
    file=open(filepath)
    for i in file:
        fields=i.replace('\n','').split('\t')
        result.append(Dish(fields[0],float(fields[1][1:]),float(fields[2])))
    return result

a=read_menu('menu3.txt')

P('\n----------part(c.3)----------\n')

def write_menu(l:list,filepath)->None:
    writefile=open(filepath,'w')
    writefile.write(str(len(l))+'\n')
    for i in l:
        s=''
        s+='{:}\t${:.2f}\t{:}\n'.format(*i)
        writefile.write(s)
    writefile.close()

write_menu(a,'menu4.txt')

P('\n----------part(d)----------\n')

Course = namedtuple('Course', 'dept num title instr units')
# Each field is a string except the number of units
ics31 = Course('ICS', '31', 'Intro to Programming', 'Kay', 4.0)
ics32 = Course('ICS', '32', 'Programming with Libraries', 'Thornton', 4.0)
wr39a = Course('Writing', '39A', 'Intro Composition', 'Alexander', 4.0)
wr39b = Course('Writing', '39B', 'Intermediate Composition', 'Gross', 4.0)
bio97 = Course('Biology', '97', 'Genetics', 'Smith', 4.0)
mgt1  = Course('Management', '1', 'Intro to Management', 'Jones', 2.0)

Student = namedtuple('Student', 'ID name level major studylist')
# All are strings except studylist, which is a list of Courses.
sW = Student('11223344', 'Anteater, Peter', 'FR', 'PSB', [ics31, wr39a, bio97, mgt1])
sX = Student('21223344', 'Anteater, Andrea', 'SO', 'CS', [ics31, wr39b, bio97, mgt1])
sY = Student('31223344', 'Programmer, Paul', 'FR', 'COG SCI', [ics32, wr39a, bio97])
sZ = Student('41223344', 'Programmer, Patsy', 'SR', 'PSB', [ics32, mgt1])

StudentBody = [sW, sX, sY, sZ]

P('\n----------part(d.1)----------\n')

def Students_at_level(l:list,s:str)->list:
    result=[]
    for i in l:
        if i.level==s:
            result.append(i)
    return result

P('\n----------part(d.2)----------\n')

def Students_in_majors(l1:list,l2:list)->list:
    result=[]
    for i in l1:
        if i.major in l2:
            result.append(i)
    return result

P('\n----------part(d.3)----------\n')

def Students_in_class(l1:list,dept:str,num:str)->list:
    result=[]
    for i in l1:
        if Student_is_enrolled(i,dept,num):
            result.append(i)
    return result

def Course_equals(c1: Course, c2: Course) -> bool:
    ''' Return True if the department and number of c1 match the department and
	     number of c2 (and False otherwise)
    '''
    return c1.dept==c2.dept and c1.num==c2.num

def Course_on_studylist(c: Course, SL: 'list of Course') -> bool:
    ''' Return True if the course c equals any course on the list SL (where equality
	     means matching department name and course number) and False otherwise.
    '''
    for i in SL:
        if Course_equals(c,i):
            return True
    return False

def Student_is_enrolled(S: Student, department: str, coursenum: str) -> bool:
    ''' Return True if the course (department and course number) is on the student's
	     studylist (and False otherwise)
    '''
    for i in S.studylist:
        if i.dept==department and i.num==coursenum:
            return True
    return False

P('\n----------part(d.4)----------\n')

def Student_names(l:list)->list:
    return [i.name for i in l]

P('\n----------part(d.5)----------\n')

def Students_average_units(l:list)->float:
    result=0
    for i in l:
        for j in i.studylist:
            result+=j.units
    return result/len(l) if len(l) else 0

P('\n----------part(d.5.1)----------\n')

L1=Students_in_majors(StudentBody, ['CS', 'CSE', 'BIM', 'INFX', 'CGS', 'SE', 'ICS'])
P(L1)
P('\n----------part(d.5.2)----------\n')

P(Student_names(L1))
P('\n----------part(d.5.3)----------\n')

P(len(L1))
P('\n----------part(d.5.4)----------\n')

L2=Students_at_level(L1,'SR')
P(Student_names(L2))
P('\n----------part(d.5.5)----------\n')

P(len(L2))
P('\n----------part(d.5.6)----------\n')

L3=Students_at_level(StudentBody, 'SR')
P(str(len(L2)*100/len(L3))+'%')
P('\n----------part(d.5.7)----------\n')

L4=Students_in_class(Students_at_level(L1,'FR'),'ICS','31')
P(len(L4))
P('\n----------part(d.5.8)----------\n')

P(Students_average_units(Students_in_class(Students_at_level(StudentBody,'FR'),'ICS','31')))


import numpy as np
import matplotlib.pyplot as plt
from scipy import stats

plt.style.use("seaborn")
import math
import csv

file = open("../../data/eth2.csv")
csvreader = csv.reader(file)

header = next(csvreader)
x = [int(header[0])]
for row in csvreader:
    x.append(int(row[0]))

def get_beta(T):
    waste = [e % T for e in x]

    return sum(waste) / I

fig, ax = plt.subplots()
n, bins, patches = ax.hist(x, 150, density=False)
mode_idx = np.argmax(n)

mode_val = bins[mode_idx]

ticket_price = mode_val * (2 / 3)

R = sum(x)
I = len(x)

AN = (
    math.floor(R / ticket_price) - math.floor(I * get_beta(ticket_price) / ticket_price)
) / I

MN = max(np.array(x) // ticket_price)

print(f"AN: {AN}")
print(f"MN: {MN}")


import argparse
import json
import re

import openai
from openai import ChatCompletion

# openai.api_base = "https://api.wandb.ai/proxy/openai/v1"
# wandb_key = os.environ["WANDB_API_KEY"]
# openai_key = os.environ['OPENAI_API_KEY']
# openai.api_key = f"{wandb_key}:{openai_key}"

# OPENAI_MODEL = "gpt-3.5-turbo"
OPENAI_MODEL = "gpt-4"
TEMPERATURE = 0.7

def parse_example_file(file) -> (str, list[dict], list[dict]):
    with open(file, "r") as f:
        cur_user_prompt = ""
        cur_answer = ""
        section = ""
        line_number = 0
        prompt = ""
        examples = []
        prompt_examples = []

        for line in f:
            if line.startswith("####"):
                # section header
                if line.startswith("#### Prompt:"):
                    section = "prompt"
                elif line.startswith("#### Prompt Examples:"):
                    section = "prompt examples"
                elif line.startswith("#### Examples:"):
                    section = "examples"
                else:
                    raise ValueError(
                        f"Malformed section header in line {line_number} in {file}"
                    )
            else:
                if section == "prompt":
                    if line.strip != "":
                        prompt += line
                elif section == "examples" or section == "prompt examples":
                    if line.startswith("### User:"):
                        if cur_user_prompt != "":
                            raise ValueError(
                                f"Malformed file: two user prompts in a row in line {line_number} in {file}"
                            )
                        # get the string after ### User:
                        cur_user_prompt = line.split(":", 2)[1].strip()

                    elif line.startswith("### Assistant:"):
                        # get part of line after the :
                        cur_answer = line.split(":", 1)[1].strip()
                        example = {"user": cur_user_prompt,
                                   "answer": cur_answer}
                        if section == "prompt examples":
                            prompt_examples.append(example)
                        else:
                            examples.append(example)
                        cur_user_prompt = ""
                    elif line.startswith("###"):
                        raise ValueError(
                            f"Malformed file: line needs to start with ### User or ### Assistant at line {line_number} in {file}"
                        )

            line_number += 1
        return (prompt, prompt_examples, examples)

def create_training_data_file(files):
    examples = []
    for file in files:
        prompt, prompt_examples, file_examples = parse_example_file(file)
        examples.append(file_examples)

    for example in examples:
        for line in example:
            print(json.dumps(line))

def prompt_examples_to_string(prompt_examples: [dict]):
    prompt_examples_str = ""
    for example in prompt_examples:
        prompt_examples_str += (
            f"### User: {example['user']}\n### Assistant: {example['answer']}\n"
        )
    return prompt_examples_str

def create_prompt(prompt: str, prompt_examples: list[dict]):
    data_collection_prompt = f"""
I am collecting training data for a voice assistant.

The voice assistant has the command:
{prompt}
Some examples of how a user might say this command and the response is:
{prompt_examples_to_string(prompt_examples)}
Generate an examples of a user querying this command and the correct response.

Use the following format
User:
Assistant:
"""

    return data_collection_prompt

format_function = [
    {
        "name": "validate",
        "description": "validates the user assistant pair",
        "parameters": {
            "type": "object",
            "properties": {
                 "user": {"type": "string", "description": "The query from the user" },
                 "assistant":{"type": "string", "description": "The command to call" }}
        },
        "required": ["user", "assistant"],
    },
]

def create_prompts(files: [str]):
    for file in files:
        prompt, prompt_examples, file_examples = parse_example_file(file)

        data_collection_prompt = create_prompt(prompt, prompt_examples)
        print(data_collection_prompt)

def collect_training_data(file: str, n_generations: int=5):
    prompt, prompt_examples, file_examples = parse_example_file(file)
    data_collection_prompt = create_prompt(prompt, prompt_examples)
    print(data_collection_prompt)

    response = ChatCompletion.create(
        model=OPENAI_MODEL,
        messages=[
            {
                "role": "system",
                "content": "You are helpful assistant generating training data for a voice assistant.",
            },
            {"role": "user", "content": data_collection_prompt},
        ],
        n = n_generations,
        functions=format_function,
        function_call = {"name": "validate"},
        temperature=TEMPERATURE,
    )
    for c in response["choices"]:
        try:
            if c.message.function_call != None:
                response = json.loads(c.message.function_call.arguments)
                print("### User: {user}\n### Assistant: {assistant}\n\n".format(**response))
        except:
            pass

def strip_ansi_codes(s):
    return re.sub(r"\x1b\[([0-9,A-Z]{1,2}(;[0-9]{1,2})?(;[0-9]{3})?)?[m|K]?", "", s)

def generate_other_examples(files):
    lines = set()  # set of lines - want to avoid duplicates
    for file in files:
        with open(file, "r") as f:
            for line in f:
                line = strip_ansi_codes(line)
                line = line.strip()
                if (
                    line == ""
                    or line.startswith("(")
                    or line.startswith("[")
                    or line.startswith("*")
                    or line.startswith(".")
                    or line.startswith("-")
                ):
                    continue
                else:
                    lines.add(line)

    for line in lines:
        print(f"### User: {line}\n### Assistant: other()\n")

if __name__ == "__main__":
    argparser = argparse.ArgumentParser(

        description='Generate training data from example files')

    argparser.add_argument('--files', metavar='file', type=str, nargs='+',
                           help='example files to parse')
    argparser.add_argument('-t', '--training-data-file', action='store_true',
                           help="generate training data file")
    argparser.add_argument('-p', '--prompt', action='store_true',
                           help="generate training data collection prompt")
    argparser.add_argument('-c', '--collect-training-data', action='store_true',
                           help="collect training data")
    argparser.add_argument('-o', '--generate-other-examples', action='store_true',
                           help="generate other examples")
    argparser.add_argument('-n', '--n_generations', type=int, default=5,
                           help="number of generations to run")

    args = argparser.parse_args()
    if args.training_data_file:
        create_training_data_file(args.files)
    elif args.prompt:
        create_prompts(args.files)
    elif args.collect_training_data:
        collect_training_data(args.files[0], n_generations=args.n_generations)
    elif args.generate_other_examples:
        generate_other_examples(args.files)


import cv2
import numpy as np

def start():
    cap = cv2.VideoCapture(0)

    while True:
        ret, frame = cap.read()

        cv2.imshow('frame', frame)
        if cv2.waitKey(1) & 0xFF == 27:
            break

    cap.release()
    cv2.destroyAllWindows()

def close(event, x, y, flags, param):
    if event == cv2.EVENT_LBUTTONDOWN:
        cv2.destroyAllWindows()
        start()

def main_page():
    img = cv2.imread('eye.jpg', 1)
    cv2.imshow('start', img)
    cv2.setMouseCallback('start', close)
    k = cv2.waitKey(0) & 0xFF

if __name__ == '__main__':
    main_page()