# -*- coding: utf-8 -*-
from __future__ import unicode_literals

from django.db import models, migrations
from django.conf import settings

class Migration(migrations.Migration):

    dependencies = [
        migrations.swappable_dependency(settings.AUTH_USER_MODEL),
    ]

    operations = [
        migrations.CreateModel(
            name='Quote',
            fields=[
                ('id', models.AutoField(verbose_name='ID', serialize=False, auto_created=True, primary_key=True)),
                ('title', models.CharField(max_length=255)),
                ('body', models.TextField()),
                ('type', models.CharField(max_length=15, choices=[(b'internal', b'internal source'), (b'external', b'external source')])),
                ('created', models.DateTimeField(auto_now_add=True)),
                ('image', models.ImageField(null=True, upload_to=b'', blank=True)),
                ('involved', models.ManyToManyField(related_name='quotes', null=True, to=settings.AUTH_USER_MODEL, blank=True)),
                ('poster', models.ForeignKey(to=settings.AUTH_USER_MODEL)),
            ],
            options={
            },
            bases=(models.Model,),
        ),
        migrations.CreateModel(
            name='Screenshot',
            fields=[
                ('id', models.AutoField(verbose_name='ID', serialize=False, auto_created=True, primary_key=True)),
            ],
            options={
            },
            bases=(models.Model,),
        ),
    ]


from text_embeddings.preprocessing.read_pdf import *
from user_interface.cli import *
import gensim
from gensim.models import Word2Vec

'''------Different word embeddings-------'''

def main(file_paths: list):
    # https://www.geeksforgeeks.org/python-word-embedding-using-word2vec/

    lower_case_tokens = []
    for path in file_paths:
        text = pdf_to_str(path)

        tokens = tokenize(text)
        lower_case_tokens.append(list(map(lambda x: x.lower(), tokens)))

    # Create CBOW model of both files
    model1 = gensim.models.Word2Vec(lower_case_tokens, min_count = 1, vector_size = 100, window = 5)

    # Print results
    print("Cosine similarity between 'bahamas' " + "and 'wealth' - CBOW : ", model1.wv.similarity('bahamas', 'wealth'))

    print("Cosine similarity between 'bahamas' " + "and 'credit' - CBOW : ", model1.wv.similarity('bahamas', 'credit'))

    # Create Skip Gram model of both files
    model2 = gensim.models.Word2Vec(lower_case_tokens, min_count = 1, vector_size = 100, window = 5, sg = 1)

    # Print results
    print("Cosine similarity between 'bahamas' " + "and 'wealth' - Skip Gram : ", model2.wv.similarity('bahamas', 'wealth'))

    print("Cosine similarity between 'bahamas' " + "and 'credit' - Skip Gram : ", model2.wv.similarity('bahamas', 'credit'))


import svgwrite
import math

GOLDEN_RATIO = 0.618
COLOR = 'blue'

def step_fractal_tree(img, iterations, color, x1, y1, angle, length):
    """
    Recursive procedure creating a line and than calls itself twice for
    making two separate branches.
    """
    x2 = x1 + length*math.sin(math.radians(angle))
    y2 = y1 + length*math.cos(math.radians(180-angle))
    img.add(img.line((x1, y1), (x2, y2), stroke=color))
    if iterations > 1:
        step_fractal_tree(img, iterations - 1, color, x2,
                          y2, angle - 30, 0.7*length)
        step_fractal_tree(img, iterations - 1, color, x2,
                          y2, angle + 30, 0.7*length)

def fractal_tree(img, iterations, size, color):
    """
    Initial step for making a fractal tree.
    """
    step_fractal_tree(img, iterations, color, 0, 0, 0, size)

def step_fractal_triangle(img, iterations, color, triangle_coordinates):
    """
    Makes one iteration of splitting the triangle into four equal triangles.
    Than it recursively continues by splitting three triangles in the corners.
    """
    new = []
    old = triangle_coordinates

    #Finds coordinates for splitting the triangle.
    for i in ([0, 1], [1, 2], [2, 0]):
        new.append([(old[i[0]][0] + old[i[1]][0])/2,
                    (old[i[0]][1] + old[i[1]][1])/2])

    #Makes three lines splitting old triangle into four equal triangles.
    for i in ([0, 1], [1, 2], [2, 0]):
        img.add(img.line(new[i[0]], new[i[1]], stroke=color))

    #Starts recursion on new triangles made in the corners of the old triangle.
    if iterations>1:
        step_fractal_triangle(img, iterations - 1, color,
                              [old[0], new[0], new[2]])
        step_fractal_triangle(img, iterations - 1, color,
                              [new[0], old[1], new[1]])
        step_fractal_triangle(img, iterations - 1, color,
                              [new[2], new[1], old[2]])

def fractal_triangle(img, iterations, size, color):
    """
    Initial step of making Sierpinski triangle. Than it activates the
    recursive procedure of splitting the triangle.
    """
    angle = math.radians(60)
    new = [(0, 0), (size*math.cos(angle), -size*math.sin(angle)), (size, 0)]

    for i in ([0, 1], [1, 2], [2, 0]):
        img.add(img.line(new[i[0]], new[i[1]], stroke=color))

    step_fractal_triangle(img, iterations - 1, color, new)

def golden_fractal_step(img, iterations, color, center, start, dir, mode):
    """
    Recursive procedure for making fractals based on golden spiral.
    """
    p = img.path('m' + str(start[0]) + ',' + str(start[1]),
                 fill = 'none', stroke = color)

    n = (center[1] - start[1], -center[0] + start[0])

    if dir==mode:
        target = (center[0] + n[0], center[1] + n[1])
    else:
        target = (center[0] - n[0], center[1] - n[1])

    next1 = (target[0] + n[0]*GOLDEN_RATIO, target[1] + n[1]*GOLDEN_RATIO)
    next2 = (target[0] - n[0]*GOLDEN_RATIO, target[1] - n[1]*GOLDEN_RATIO)

    p.push_arc(target, r=math.fabs(n[0] + n[1]), rotation=0,
               large_arc=False, absolute=True, angle_dir=dir)
    img.add(p)

    if iterations>1:
        golden_fractal_step(img, iterations - 1, color,
                            next1, target, '-', mode)
        golden_fractal_step(img, iterations - 1, color,
                            next2, target, '+', mode)

def golden_fractal_imp(img, iterations, size, color):
    """
    Initial step for making a fractal based on golden spiral.
    """
    golden_fractal_step(img, iterations - 1, color,
                        (-size,0), (0, 0), '-', '-')
    golden_fractal_step(img, iterations - 1, color,
                        (size,0), (0, 0), '+', '-',)

def golden_fractal_exp(img, iterations, size, color):
    """
    Initial step for making an expansive fractal based on golden spiral.
    """
    golden_fractal_step(img, iterations - 1, color,
                        (-size,0), (0, 0), '-', '+')
    golden_fractal_step(img, iterations - 1, color,
                        (size,0), (0, 0), '+', '+')

def make_fractal(fname, function, iterations, base_size, color):
    """
    General function for creating fractal based on the given function.
    """
    img = svgwrite.Drawing(fname, profile='full')
    function(img, iterations, base_size, color)
    img.save()

#make_fractal("ftree.svg", fractal_tree, 12, 300, COLOR)
#make_fractal("ftriangle.svg", fractal_triangle, 7, 300, COLOR)
#make_fractal('goldenf1.svg', golden_fractal_imp, 10, 300, COLOR)
#make_fractal('goldenf2.svg', golden_fractal_exp, 10, 300, COLOR)


import os
import urllib.request

from tqdm import tqdm
from zipfile import ZipFile
import shutil

class DownloadProgressBar(tqdm):
    """Progress bar wrapper for downloading"""

    def update_to(self, b=1, bsize=1, tsize=None):
        """Update progress bar"""
        if tsize is not None:
            self.total = tsize
        self.update(b * bsize - self.n)

def download_url(url, output_path):
    """https://stackoverflow.com/a/53877507"""
    with DownloadProgressBar(unit='B', unit_scale=True,
                             miniters=1, desc=url.split('/')[-1]) as bar:
        urllib.request.urlretrieve(url, filename=output_path, reporthook=bar.update_to)

def download_cldr_data(download_path='temp/cldr_data.zip', version='latest'):
    """Download the latest CLDR data from the website"""
    # The zip of very early versions is called cldr.zip not core.zip
    versions_where_zip_is_cldr = ['1.2.0', '1.1.1', '1.1.0']
    zip_name = 'core' if version not in versions_where_zip_is_cldr else 'cldr'
    download_link = f"https://unicode.org/Public/cldr/{version}/{zip_name}.zip"
    os.makedirs('temp', exist_ok=True)
    download_url(download_link, download_path)

def extract_cldr_data(download_path='temp/cldr_data.zip', version='latest'):
    """Extract CLDR data from downloaded zip, only part of the file needs unzipping"""
    # In early versions the main folder is in a different location
    versions_with_no_common = ['1.6.1', '1.6.0', '1.5.1', '1.5.0', '1.4.1', '1.4.0',
                               '1.2.0', '1.1.1', '1.1.0']
    main_path = 'common/' if version not in versions_with_no_common else ''

    with ZipFile(download_path, 'r') as zip_obj:
        list_of_file_names = zip_obj.namelist()
        file_names = list(filter(lambda x: f'{main_path}main/' in x, list_of_file_names))
        zip_obj.extractall(members=file_names, path='temp')
    # Just to make it a little nicer I will move the files
    shutil.rmtree('temp/cldr_data', ignore_errors=True)
    shutil.copytree(f'temp/{main_path}main/', 'temp/cldr_data')
    shutil.rmtree(f'temp/{main_path}main/')

def download_and_extract_cldr_data(download_path='temp/cldr_data.zip', version='latest'):
    """Download latest CLDR data and extract it"""
    download_cldr_data(download_path, version)
    extract_cldr_data(download_path, version)

if __name__ == "__main__":
    download_and_extract_cldr_data()


import networkx as nx
import matplotlib.pyplot as plt
import numpy

s = 4
G = nx.grid_graph(dim=[5,s])

s=(0,0) ; t=(0,4)
outs = [(2,2)]
G.remove_nodes_from( outs )

print(f" G.nodes() ={ G.nodes }")
#outs = (1,2)
#G.remove_node( outs )

pos=nx.spectral_layout(G)
pos=nx.kamada_kawai_layout(G)

ncon= nx.node_connectivity(G)
print("> Node connectivity of G=", ncon )
ncon= nx.edge_connectivity(G)
print("> Edge connectivity of G=", ncon )

avgcon= nx.average_node_connectivity(G)
print("> Average connectivity of G=", avgcon )

cutedge = nx.minimum_edge_cut(G)
print("> Minimum edge cut set of G=", cutedge )

cutedge = nx.minimum_edge_cut(G, s=(1,2), t=(2,1))
for i,w in enumerate(cutedge):
    print(f">  {i}: edge cut between s(1,2) and t(2,1) = {w}" )

# Drawing Code
nx.draw_networkx_nodes(G, pos, node_size= 1200, node_color='b', alpha=0.5)
nx.draw_networkx_labels(G,pos,font_size=11, font_color='white', font_family='sans-serif')
nx.draw_networkx_edges(G, pos, width=2)

plt.axis('off')
plt.show()


import argparse
import datetime
import json
import os
import time
import tqdm

import numpy as np
import pandas as pd
import sklearn.metrics as sm

import torch
import torch.nn as nn
from torch.utils.data import DataLoader

from atom3d.datasets import LMDBDataset
from scipy.stats import spearmanr

from model import CNN3D_PPI
from data import CNN3D_Dataset

def major_vote(results):
    data = []
    for ensemble, df in results.groupby('id'):
        true = int(df['true'].unique()[0])
        num_zeros = np.sum(df['pred'] == 0)
        num_ones = np.sum(df['pred'] == 1)
        majority_pred = int(df['pred'].mode().values[0])
        avg_prob = df['prob'].astype(np.float).mean()
        data.append([ensemble, true, majority_pred, avg_prob, num_zeros, num_ones])
    vote_df = pd.DataFrame(data, columns=['id', 'true', 'pred', 'avg_prob',
                                          'num_pred_zeros', 'num_pred_ones'])
    return vote_df

def compute_stats(df):
    results = major_vote(df)
    res = {}
    all_true = results['true'].astype(np.int8)
    all_pred = results['pred'].astype(np.int8)
    res['auroc'] = sm.roc_auc_score(all_true, all_pred)
    res['auprc'] = sm.average_precision_score(all_true, all_pred)
    res['acc'] = sm.accuracy_score(all_true, all_pred.round())
    res['bal_acc'] = \
        sm.balanced_accuracy_score(all_true, all_pred.round())
    return res

# Construct model
def conv_model(in_channels, spatial_size, args):
    num_conv = args.num_conv
    conv_filters = [32 * (2**n) for n in range(num_conv)]
    conv_kernel_size = 3
    max_pool_positions = [0, 1]*int((num_conv+1)/2)
    max_pool_sizes = [2]*num_conv
    max_pool_strides = [2]*num_conv
    fc_units = [512]
    top_fc_units = [512]*args.num_final_fc_layers

    model = CNN3D_PPI(
        in_channels, spatial_size,
        args.conv_drop_rate,
        args.fc_drop_rate,
        args.top_nn_drop_rate,
        conv_filters, conv_kernel_size,
        max_pool_positions,
        max_pool_sizes, max_pool_strides,
        fc_units,
        top_fc_units,
        batch_norm=args.batch_norm,
        dropout=not args.no_dropout)
    return model

def train_loop(model, loader, criterion, optimizer, device):
    model.train()

    loss_all = 0
    total = 0
    progress_format = 'train loss: {:6.6f}'
    with tqdm.tqdm(total=len(loader), desc=progress_format.format(0)) as t:
        for i, data in enumerate(loader):
            feature_left = data['feature_left'].to(device).to(torch.float32)
            feature_right = data['feature_right'].to(device).to(torch.float32)
            labels = data['label'].to(device).to(torch.float32)
            # zero the parameter gradients
            optimizer.zero_grad()
            # forward + backward + optimize
            output = model(feature_left, feature_right)
            loss = criterion(output, labels)
            loss.backward()
            loss_all += loss.item() * len(labels)
            total += len(labels)
            optimizer.step()
            # stats
            t.set_description(progress_format.format(np.sqrt(loss_all/total)))
            t.update(1)

    return np.sqrt(loss_all / total)

@torch.no_grad()
def test(model, loader, criterion, device):
    model.eval()

    loss_all = 0
    total = 0

    ids = []
    y_true = []
    y_probs = []
    y_pred = []

    for data in loader:
        feature_left = data['feature_left'].to(device).to(torch.float32)
        feature_right = data['feature_right'].to(device).to(torch.float32)
        labels = data['label'].to(device).to(torch.float32)
        output = model(feature_left, feature_right)

        loss = criterion(output, labels)
        loss_all += loss.item() * len(labels)
        total += len(labels)

        ids.extend(data['id'])
        y_true.extend(labels.int().tolist())
        y_probs.extend(output.tolist())

        preds = torch.round(output).int()
        y_pred.extend(preds.tolist())

    results_df = pd.DataFrame(
        np.array([ids, y_true, y_probs, y_pred]).T,
        columns=['id', 'true', 'prob', 'pred'],
        )

    stats = compute_stats(results_df)
    return np.sqrt(loss_all / total), stats, results_df

def save_weights(model, weight_dir):
    torch.save(model.state_dict(), weight_dir)

def train(args, device, test_mode=False):
    print("Training model with config:")
    print(str(json.dumps(args.__dict__, indent=4)) + "\n")

    # Save config
    with open(os.path.join(args.output_dir, 'config.json'), 'w') as f:
        json.dump(args.__dict__, f, indent=4)

    np.random.seed(args.random_seed)
    torch.manual_seed(args.random_seed)

    train_dataset = CNN3D_Dataset(
        os.path.join(args.data_dir, 'train'),
        testing=False,
        random_seed=args.random_seed,
        )
    val_dataset = CNN3D_Dataset(
        os.path.join(args.data_dir, 'val'),
        testing=False,
        random_seed=args.random_seed,
        )
    test_dataset = CNN3D_Dataset(
        os.path.join(args.data_dir, 'test'),
        testing=True,
        random_seed=args.random_seed,
        )

    train_loader = DataLoader(train_dataset, args.batch_size)
    val_loader = DataLoader(val_dataset, args.batch_size)
    test_loader = DataLoader(test_dataset, args.batch_size)

    for data in train_loader:
        in_channels, spatial_size = data['feature_left'].size()[1:3]
        print('num channels: {:}, spatial size: {:}'.format(in_channels, spatial_size))
        break

    model = conv_model(in_channels, spatial_size, args)
    print(model)
    model.to(device)

    prev_val_loss = np.Inf
    best_val_loss = np.Inf
    best_val_auroc = 0
    best_stats = None

    criterion = nn.BCELoss()
    criterion.to(device)
    optimizer = torch.optim.Adam(model.parameters(), lr=args.learning_rate)

    for epoch in range(1, args.num_epochs+1):
        start = time.time()
        train_loss = train_loop(model, train_loader, criterion, optimizer, device)
        val_loss, stats, val_df = test(model, val_loader, criterion, device)
        elapsed = (time.time() - start)
        print(f'Epoch {epoch:03d} finished in : {elapsed:.3f} s')
        print(f"\tTrain loss {train_loss:.4f}, Val loss: {val_loss:.4f}, "
              f"Val AUROC: {stats['auroc']:.4f}, Val AUPRC: {stats['auroc']:.4f}")
        #if stats['auroc'] > best_val_auroc:
        if val_loss < best_val_loss:
            print(f"\nSave model at epoch {epoch:03d}, val_loss: {val_loss:.4f}, "
                  f"auroc: {stats['auroc']:.4f}, auprc: {stats['auroc']:.4f}")
            save_weights(model, os.path.join(args.output_dir, f'best_weights.pt'))
            best_val_loss = val_loss
            best_val_auroc = stats['auroc']
            best_stats = stats
        if args.early_stopping and val_loss >= prev_val_loss:
            print(f"Validation loss stopped decreasing, stopping at epoch {epoch:03d}...")
            break
        prev_val_loss = val_loss

    if test_mode:
        model.load_state_dict(torch.load(os.path.join(args.output_dir, f'best_weights.pt')))
        test_loss, stats, test_df = test(model, test_loader, criterion, device)
        test_df.to_pickle(os.path.join(args.output_dir, 'test_results.pkl'))
        print(f"Test loss: {test_loss:.4f}, Test AUROC: {stats['auroc']:.4f}, "
              f"Test AUPRC: {stats['auprc']:.4f}")
        test_file = os.path.join(args.output_dir, f'test_results.txt')
        with open(test_file, 'w') as f:
            f.write(f"test_loss\tAUROC\tAUPRC\n")
            f.write(f"{test_loss:}\t{stats['auroc']:}\t{stats['auprc']:}\n")

if __name__=="__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--data_dir', type=str, default=os.environ['PPI_DATA'])
    parser.add_argument('--mode', type=str, default='test',
                        choices=['train', 'test'])
    parser.add_argument('--output_dir', type=str, default=os.environ['LOG_DIR'])
    parser.add_argument('--unobserved', action='store_true', default=False)

    parser.add_argument('--learning_rate', type=float, default=0.00001)
    parser.add_argument('--conv_drop_rate', type=float, default=0.1)
    parser.add_argument('--fc_drop_rate', type=float, default=0.25)
    parser.add_argument('--top_nn_drop_rate', type=float, default=0.25)
    parser.add_argument('--num_epochs', type=int, default=30)

    parser.add_argument('--num_conv', type=int, default=4)
    parser.add_argument('--num_final_fc_layers', type=int, default=2)
    parser.add_argument('--batch_norm', action='store_true', default=False)
    parser.add_argument('--no_dropout', action='store_true', default=False)

    parser.add_argument('--batch_size', type=int, default=16)
    parser.add_argument('--early_stopping', action='store_true', default=False)
    parser.add_argument('--random_seed', type=int, default=int(np.random.randint(1, 10e6)))

    args = parser.parse_args()

    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    # Set up output dir
    args.output_dir = os.path.join(args.output_dir, 'ppi')
    assert args.output_dir != None
    if args.unobserved:
        args.output_dir = os.path.join(args.output_dir, 'None')
        os.makedirs(args.output_dir, exist_ok=True)
    else:
        num = 0
        while True:
            dirpath = os.path.join(args.output_dir, str(num))
            if os.path.exists(dirpath):
                num += 1
            else:
                args.output_dir = dirpath
                print('Creating output directory {:}'.format(args.output_dir))
                os.makedirs(args.output_dir)
                break

    print(f"Running mode {args.mode:} with seed {args.random_seed:} "
          f"and output dir {args.output_dir}")
    train(args, device, args.mode=='test')


import nltk
import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt
from scipy import stats
import numpy as np
from nltk.tag import pos_tag
from nltk import word_tokenize
from nltk.corpus import stopwords
import string
from collections import defaultdict
from collections import Counter
from nltk.stem import WordNetLemmatizer
from sklearn.metrics.pairwise import cosine_similarity
import ast
from statannot import add_stat_annotation
import copy
from statistics import mean, stdev
from math import sqrt
#import eval_lib as elib
import linguistic_stats_functions as elib
import sys

stop = stopwords.words('english') + list(string.punctuation)

## these variables are the same as the ones in the last two tables of appendix G
compression = ['text_len_content', 'content_prop',
                          'cont_noun_prop', 'cont_adj_prop', 'cont_verb_prop']
givenness = ['prop_the', 'prop_a',
                        'prop_seen_markers',
                        'prop_givenness_markers']
reuse = ['prop_content_from_prev', 'prop_reused_content_bigrams',
                 'prop_nouns_in_reused_content_tokens', 'prop_adjs_in_reused_content_tokens', 'prop_verbs_in_reused_content_tokens',
                 'prop_nn_content_bigrams_reused']

### functions
def get_prop_tag_content(tag, taglist, wdlist):
    tags = 0
    denom = 0
    for i in range(len(taglist)):
        if wdlist[i] not in stop: ## only content
            denom += 1
            if taglist[i] == tag:
                tags += 1
    return np.NaN if denom < 1 else tags/denom

if len(sys.argv) == 1:
    print('specify files to process as a list of command line arguments')
    sys.exit()

## Create dataframe from output files
print('loading files...')
files = sys.argv[1:]
#files = ['all_copy_TEST.txt', 'all_base_TEST.txt', 'all_histatt_TEST.txt']
print(files)
res = elib.get_dataframe_from_files(files)
#res.to_csv('model_output_summary.csv')

## Calculate stats
print('tagging...')
res['tags'] = res.apply(lambda row: [item[1] for item in pos_tag(row.text, tagset='universal')], axis = 1)
res['prev_tokens'] = res.apply(lambda row: [item[1] for item in pos_tag(row.prev_utt_full, tagset='universal')], axis = 1)

res['content_text'] = res.apply(lambda row: [i for i in row.text if i not in stop], axis = 1)
res['content_tokens'] = res.apply(lambda row: elib.content_tag_list(row.text, row.tags), axis = 1)
res['content_prev'] = res.apply(lambda row: [i for i in row.prev_utt_full if i not in stop], axis = 1)
res['content_tokens_prev'] = res.apply(lambda row: elib.content_token_list(row.prev_utt_full, row.prev_tokens), axis = 1)

res['reused_wds'] = res.apply(lambda row:  elib.get_shared_words(row.text, row.prev_utt_full), axis = 1)
res['reused_tags'] = res.apply(lambda row:  elib.get_shared_wd_tags(row.text, row.prev_utt_full, row.tags), axis = 1)

print('compression stats...')
res['text_len_content'] = res.apply(lambda row: len(row.content_text), axis = 1)
res['content_prop'] = res.apply(lambda row: elib.content_prop_text(row.text), axis = 1)

res['temp'] = res.apply(lambda row: elib.content_token_prop(row.text), axis = 1)
res['cont_adj_prop'] = res.apply(lambda row: row.temp[0], axis = 1)
res['cont_verb_prop'] = res.apply(lambda row: row.temp[1], axis = 1)
res['cont_noun_prop'] = res.apply(lambda row: row.temp[2], axis = 1)

print('givenness stats...')
res['prop_givenness_markers'] = res.apply(lambda row: elib.prop_givenness(row.text), axis = 1)
res['prop_the'] = res.apply(lambda row: elib.prop_the(row.text), axis = 1)
res['prop_a'] = res.apply(lambda row: elib.prop_a(row.text), axis = 1)
res['prop_seen_markers'] = res.apply(lambda row: elib.prop_seen_markers(row.text), axis = 1)

print('reuse stats...')

res['prop_content_from_prev'] = res.apply(lambda row:   elib.prop_content_reused(row.text, row.prev_utt_full), axis = 1)

res['content_bigrams'] = res.apply(lambda row: elib.get_content_bigram_strings(row.text), axis = 1)
res['content_bigrams_prev'] = res.apply(lambda row: elib.get_content_bigram_strings(row.prev_utt_full), axis = 1)
res['count_content_bigrams_reused'] = res.apply(lambda row: elib.count_words_reused(row.content_bigrams_prev, row.content_bigrams), axis = 1)
res['prop_reused_content_bigrams'] = res.apply(lambda row: row.count_content_bigrams_reused/float(max(1, len(row.content_bigrams))), axis = 1)

res['prop_nouns_in_reused_content_tokens'] = res.apply(lambda row:  elib.get_prop_tag_content('NOUN', row.reused_tags, row.reused_wds), axis = 1)
res['prop_adjs_in_reused_content_tokens'] = res.apply(lambda row:  elib.get_prop_tag_content('ADJ', row.reused_tags, row.reused_wds), axis = 1)
res['prop_verbs_in_reused_content_tokens'] = res.apply(lambda row:  elib.get_prop_tag_content('VERB', row.reused_tags, row.reused_wds), axis = 1)

res['prop_nn_content_bigrams_reused'] = res.apply(lambda row:   elib.prop_nn_bigrams_reused(row.text, row.tags, row.prev_utt_full), axis = 1)

#res.to_csv('model_output_summary_stats.csv')
res['mention'] = res.first_mention

## firstly calculate the compression and givenness stats: first vs the later are compared
rows = []
cols = ['comparison', 'model', 'version', 'meanFirst', 'meanLater', 'cohensD', 'sig', 'pval', 'tstat']
aspects = compression + givenness
for aspect in aspects:
    for version in res.version.unique():
        row = elib.compare(res, 'first', 'later', version, aspect, aspect)
        rows.append(row)
first_later = pd.DataFrame(rows, columns=cols)
first_later.to_csv('firstVlater_stats_out.csv')

## now reuse: only the later utterances are inspected: thus we test differences between human and model
cols = ['comparison', 'model', 'version', 'meanHuman', 'meanGenerated', 'cohensD', 'sig', 'pval', 'tstat', 'mention']
rows = []
versions = res.version.unique()
versions = [i for i in versions if i not in ['origional']]
human = 'origional'
for col_name in reuse:
    mention = 'later'
    for version in versions:
        row = elib.get_rows(res, version, mention, human, col_name)
        rows.append(row)
reuse = pd.DataFrame(rows, columns=cols)
reuse.to_csv('reuse_stats_out.csv')


import numpy as np

import matplotlib.pyplot as plt
from matplotlib.ticker import MultipleLocator, AutoMinorLocator
from matplotlib.colors import Normalize, LogNorm

import tensorflow as tf

from argparse import ArgumentParser
from pathlib import Path
from tqdm import tqdm
import h5py
import os
import os.path

from xp_utils import XPSampler, sqrt_icov_eigen, calc_invs_eigen
from utils import batch_apply_tf
from model import GaussianMixtureModel, FluxModel, chi_band, gaussian_prior, \
                  grads_stellar_model, grads_stellar_params, \
                  get_batch_iterator,  identify_outlier_stars, \
                  identify_flux_outliers, train_stellar_model, \
                  plot_gmm_prior, assign_variable_padded, corr_matrix, \
                  calc_stellar_fisher_hessian, load_data, save_as_h5, load_h5
import model
import plot_utils

nscale = 1.

def load_training_data(fname, validation_frac=0.2, seed=1, thin=1):
    # Load training data
    with h5py.File(fname, 'r') as f:
        d = {key:f[key][:][::thin] for key in f.keys()}
        sample_wavelengths = f['flux'].attrs['sample_wavelengths'][:]
    #sample_wavelengths = np.load('wl.npy').astype('f4')

    # Ensure that certain fields are in float32 format
    f4_keys = [
        'flux', 'flux_err', 'flux_sqrticov',
        'flux_cov_eival_max', 'flux_cov_eival_min',
        'plx', 'plx_err',
        'stellar_ext', 'stellar_ext_err',
        'stellar_type', 'stellar_type_err'
    ]
    for k in f4_keys:
        d[k] = d[k].astype('f4')

    # Warn about NaNs
    check_nan_keys = [
        'flux', 'flux_err', 'flux_sqrticov',
        'flux_cov_eival_max', 'flux_cov_eival_min',
        'plx', 'plx_err',
        'stellar_type', 'stellar_type_err'
    ]
    for k in check_nan_keys:
        if np.any(np.isnan(d[k])):
            raise ValueError(f'NaNs detected in {k}!')

    # Replace NaN extinctions with 0 +- infinity
    idx = np.isnan(d['stellar_ext']) | np.isnan(d['stellar_ext_err'])
    print(f'Replacing {np.count_nonzero(idx)} NaN extinctions with 0+-inf.')
    d['stellar_ext'][idx] = 0.
    d['stellar_ext_err'][idx] = np.inf

    # Shuffle data, and put last X% into validation set
    rng = np.random.default_rng(seed=seed)
    n = len(d['gdr3_source_id'])
    shuffle_idx = np.arange(n, dtype='i8')
    rng.shuffle(shuffle_idx)

    n_val = int(np.ceil(n * validation_frac))
    d_train = {'shuffle_idx': shuffle_idx[:-n_val]}
    d_val = {'shuffle_idx': shuffle_idx[-n_val:]}

    # Copy data into training and validation sets
    for key in d:
        d_train[key] = d[key][d_train['shuffle_idx']]
        d_val[key] = d[key][d_val['shuffle_idx']]

    return d_train, d_val, sample_wavelengths

def down_sample_weighing(x_ini, all_x, bin_edges,
                         n_bins=100, max_upsampling=5.):
    # Use high-Extinction stars for empirical distribution of xi
    bin_edges = np.hstack([[-np.inf], bin_edges, [np.inf]])

    # Calculate the emperical distribution of x_ini under the given bins
    bin_indices = np.digitize(x_ini, bin_edges)
    counts = np.bincount(bin_indices, minlength=n_bins+3)
    weights = counts / counts.sum()  # convert counts to probabilities
    weights_per_bin = 1./(weights+0.05)

    weights_per_bin[0] = 0
    weights_per_bin[-1] = 0

    # Weigh all stars by the inverse of density of the ini sample
    bin_indices_all = np.digitize(all_x, bin_edges)
    weights_per_star = weights_per_bin[bin_indices_all]

    # Normalize the weights per star by median value
    weights_per_star /= (1e-5 + np.median(weights_per_star))

    weights_per_star = soft_clip_weights(weights_per_star, max_upsampling)

    return weights_per_star.astype('f4')

def plot_param_histograms_1d(stellar_type, weights, title, fname):
    n_params = stellar_type.shape[1]
    figsize = (3*n_params,3)

    fig,ax_arr = plt.subplots(
        1, n_params,
        figsize=figsize,
        layout='constrained'
    )

    for i,(ax,p) in enumerate(zip(ax_arr.flat,stellar_type.T)):
        p_range = (np.min(p), np.max(p))

        kw = dict(
            range=p_range, bins=100,
            log=True, density=True,
            histtype='step',
            alpha=0.8
        )
        ax.hist(p, label=r'$\mathrm{raw}$', **kw)
        ax.hist(p, label=r'$\mathrm{weighted}$', weights=weights, **kw)

        ax.xaxis.set_minor_locator(AutoMinorLocator())
        ax.grid(True, which='major', alpha=0.2)
        ax.grid(True, which='minor', alpha=0.05)

        ax.set_xlabel(rf'$\mathrm{{parameter\ {i}}}$')
        if i == 0:
            ax.legend()

    fig.suptitle(title)

    fig.savefig(fname)
    plt.close(fig)

def soft_clip_weights(weights, max_upsampling):
    return weights / (1 + weights/max_upsampling)

def weigh_prior(stellar_type_prior, d_train, scale=1., max_upsampling=5.):
    all_ln_prior = batch_apply_tf(
        stellar_type_prior.ln_prob,
        1024,
        d_train['stellar_type'],
        function=True,
        progress=True,
        numpy=True
    )
    all_prior = np.exp(scale*all_ln_prior)
    all_prior /= np.max(all_prior)
    weight_per_star = 1 / (all_prior + 1e-9)

    # Weigh stars for better representation
    weight_per_star = soft_clip_weights(weight_per_star, max_upsampling)

    return weight_per_star.astype('f4')

def train(data_fname, output_dir, stage=0, thin=1, E_low=0.1, n_epochs=512):
    '''
    Stage 0: Train the stellar model, using universal extinction curve.
    Stage 1: Train the extinction model with hqlE stars, using initial guess of
                    the slope of extinction curve.
    Stage 2: Train both extinction model and stellar model with hqlE stars
    Stage 3: Self-cleaning & Further optimize the model
    '''

    # General training parameters
    batch_size = 512
    n_bins = 100

    loss_hist = []

    # Ensure that output directories exist
    path_list = [
        'data', 'plots',
        'models/prior', 'models/flux',
        'index', 'hist_loss'
    ]
    for p in path_list:
        Path(os.path.join(output_dir,p)).mkdir(exist_ok=True, parents=True)

    # Helper function: save an h5 file, appending a base directory to filename
    def full_fn(fn):
        return os.path.join(output_dir, fn)

    if stage == 0:
        # Stage 0, begin without initial stellar model
        print(f'Loading training data from {data_fname} ...')
        d_train, d_val, sample_wavelengths = load_training_data(
            data_fname,
            thin=thin,
            seed=0,
        )
        n_train, n_val = [len(d['plx']) for d in (d_train,d_val)]
        print(f'Loaded {n_train} ({n_val}) training (validation) sources.')

        plot_training_data(d_train, output_dir=output_dir)

        # Initial guess of xi
        d_train['xi'] = np.zeros(n_train, dtype='f4')
        d_val['xi'] = np.zeros(n_val, dtype='f4')

        save_as_h5(d_val, full_fn('data/d_val.h5'))
        save_as_h5(d_train, full_fn('data/d_train.h5'))

        # Generate GMM prior
        fn_prior = full_fn('models/prior/gmm_prior')
        if os.path.exists(fn_prior+'-1.index'):
            print('Loading existing Gaussian Mixture Model prior ...')
            stellar_type_prior = GaussianMixtureModel.load(fn_prior+'-1')
        else:
            print('Generating Gaussian Mixture Model prior on stellar type ...')
            stellar_type_prior = GaussianMixtureModel(3, n_components=16)
            # Use at most this many stars to learn GMM prior
            n_prior_max = 1024*256
            stellar_type_prior.fit(d_train['stellar_type'][:n_prior_max])
            stellar_type_prior.save(fn_prior)
            print('  -> Plotting prior ...')
            plot_gmm_prior(stellar_type_prior, base_path=output_dir)

        weights_per_star = weigh_prior(stellar_type_prior, d_train)

        print('Plotting stellar-type histograms ...')
        plot_param_histograms_1d(
            d_train['stellar_type'],
            weights_per_star,
            r'$\mathrm{Training\ distribution:\ stellar\ type}$',
            os.path.join(output_dir, 'plots/training_stellar_type_hist1d')
        )

        # Generate tracks through stellar parameter space
        tracks_fn = full_fn('models/prior/tracks.h5')
        if os.path.exists(tracks_fn):
            print('Loading tracks through stellar parameter space ...')
            atm_tracks = model.load_stellar_type_tracks(tracks_fn)
        else:
            print('Generating tracks through stellar parameter space ...')
            atm_tracks = model.calculate_stellar_type_tracks(stellar_type_prior)
            model.save_stellar_type_tracks(atm_tracks, tracks_fn)
        print('Plotting tracks through stellar parameter space ...')
        for track in atm_tracks:
            plot_gmm_prior(
                stellar_type_prior,
                base_path=output_dir,
                overlay_track=track
            )

        # Initialize the parameter estimates at their measured (input) values
        for key in ('stellar_type', 'xi', 'stellar_ext', 'plx'):
            d_train[f'{key}_est'] = d_train[key].copy()
            d_val[f'{key}_est'] = d_train[key].copy()

        print('Creating flux model ...')
        n_stellar_params = d_train['stellar_type'].shape[1]
        p_low,p_high = np.percentile(
            d_train['stellar_type'],
            [16.,84.],
            axis=0
        )

        print('Training flux model on high-quality data ...')
        # Select a subset of "high-quality" stars with good measurements
        idx_hq = np.where(
            (d_train['plx']/d_train['plx_err'] > 10.)
          & np.all(d_train['stellar_type_err']<0.2, axis=1) # Type uncertainty
          & (d_train['stellar_ext_err'] < 0.1)
        )[0]
        n_hq = len(idx_hq)
        pct_hq = n_hq / d_train['plx'].shape[0]
        print(f'Training on {n_hq} ({pct_hq:.3%}) high-quality stars ...')

        stellar_model = FluxModel(
            sample_wavelengths, n_input=n_stellar_params,
            input_zp=np.median(d_train['stellar_type'],axis=0),
            input_scale=0.5*(p_high-p_low),
            hidden_size=32,
            l2=5., l2_ext_curve=0.5
        )

        # First, train the model with stars with good measurements,
        # with fixed slope of ext_curve
        title = (
            r'$\mathrm{Training\ distribution'
            r'\ (step\ 0a:\ HQ):'
            r'\ stellar\ type}$'
        )
        plot_param_histograms_1d(
            d_train['stellar_type'][idx_hq],
            weights_per_star[idx_hq],
            title,
            full_fn('plots/training_stellar_type_hist1d_step0a')
        )
        ret = train_stellar_model(
            stellar_model,
            d_train,
            weights_per_star,
            idx_train=idx_hq,
            optimize_stellar_model=True,
            optimize_stellar_params=False,
            lr_model_init=1e-4,
            batch_size=batch_size,
            n_epochs=n_epochs,
            model_update=['stellar_model','ext_curve_b'],
        )
        save_as_h5(ret, full_fn('hist_loss/hist_step0a.h5'))
        loss_hist.append(ret)
        stellar_model.save(full_fn('models/flux/xp_spectrum_model_initial'))

        # Plot stellar model along tracks through parameter space
        print('Plotting stellar model ...')
        for i,track in enumerate(atm_tracks):
            fig,ax = model.plot_stellar_model(stellar_model, track)
            fig.savefig(full_fn(f'plots/stellar_model_step0a_track{i}'))
            plt.close(fig)

        fig,_ = plot_roughness(d_train['stellar_type'], stellar_model)
        fig.savefig(full_fn('plots/roughness_step0a'))
        plt.close(fig)

        fig,_ = model.plot_extinction_curve(stellar_model, show_variation=False)
        fig.savefig(full_fn('plots/extinction_curve_step0a'))
        plt.close(fig)

        # Next, simultaneously train the stellar model and update stellar
        # parameters, using only the HQ data
        print('Training flux model and optimizing high-quality stars ...')
        ret = train_stellar_model(
            stellar_model,
            d_train,
            weights_per_star,
            idx_train=idx_hq,
            optimize_stellar_model=True,
            optimize_stellar_params=True,
            lr_model_init=1e-5*nscale,
            lr_stars_init=1e-4*nscale,
            batch_size=batch_size,
            n_epochs=n_epochs,
            model_update=['stellar_model','ext_curve_b'],
            var_update = ['atm','E','plx'],
        )
        save_as_h5(ret, full_fn('hist_loss/hist_step0b.h5'))
        loss_hist.append(ret)
        #plot_loss(ret, suffix='_intermediate')
        stellar_model.save(
            full_fn('models/flux/xp_spectrum_model_intermediate')
        )

        print('Plotting stellar model ...')
        for i,track in enumerate(atm_tracks):
            fig,ax = model.plot_stellar_model(stellar_model, track)
            fig.savefig(full_fn(f'plots/stellar_model_step0b_track{i}'))
            plt.close(fig)

        fig,_ = plot_roughness(d_train['stellar_type'], stellar_model)
        fig.savefig(full_fn('plots/roughness_step0b'))
        plt.close(fig)

        fig,_ = model.plot_extinction_curve(stellar_model, show_variation=False)
        fig.savefig(full_fn('plots/extinction_curve_step0b'))
        plt.close(fig)

        stellar_model = FluxModel.load(
            full_fn('models/flux/xp_spectrum_model_intermediate-1')
        )

        # Update E for stars with large E uncertainties (often, +-inf)
        print('Optimizing uncertain E estimates ...')
        idx_large_ext_err = np.where(d_train['stellar_ext_err'] > 0.1)[0]
        ret = train_stellar_model(
            stellar_model,
            d_train,
            weights_per_star,
            idx_train=idx_large_ext_err,
            optimize_stellar_model=False,
            optimize_stellar_params=True,
            batch_size=batch_size,
            lr_stars_init=1e-2,
            n_epochs=32*n_epochs,
            var_update=['E'],
        )

        # Next, update parameters of all the stars, holding the model fixed
        print('Optimizing all stars ...')
        ret = train_stellar_model(
            stellar_model,
            d_train,
            weights_per_star,
            optimize_stellar_model=False,
            optimize_stellar_params=True,
            batch_size=batch_size,
            n_epochs=n_epochs,
            var_update=['atm','E','plx'],
        )
        loss_hist.append(ret)

        # Self-cleaning: Identify outlier stars to exclude from further training,
        # using distance from priors
        idx_params_good = identify_outlier_stars(d_train)
        pct_good = np.mean(idx_params_good)
        print(f'Parameter outliers: {1-pct_good:.3%} of sources.')

        idx_flux_good = identify_flux_outliers(
            d_train, stellar_model,
            chi2_dof_clip=10.,
            #chi_indiv_clip=20.
        )

        pct_good = np.mean(idx_flux_good)
        print(f'Flux outliers: {1-pct_good:.3%} of sources.')

        idx_good = idx_params_good & idx_flux_good
        pct_good = np.mean(idx_good)
        print(f'Combined outliers: {1-pct_good:.3%} of sources.')

        title = (
            r'$\mathrm{Training\ distribution'
            r'\ (step\ 0c:\ cut\ flux/param\ outliers):'
            r'\ stellar\ type}$'
        )
        plot_param_histograms_1d(
            d_train['stellar_type'][idx_good],
            weights_per_star[idx_good],
            title,
            full_fn('plots/training_stellar_type_hist1d_step0c')
        )

        # Finally, simultaneously train the stellar model and update stellar
        # parameters, using all the (non-outlier) data
        print('Training flux model and optimizing all non-outlier stars ...')
        ret = train_stellar_model(
            stellar_model,
            d_train,
            weights_per_star,
            idx_train=np.where(idx_good)[0],
            optimize_stellar_model=True,
            optimize_stellar_params=True,
            lr_model_init=1e-5*nscale,
            lr_stars_init=1e-5*nscale,
            batch_size=batch_size,
            n_epochs=n_epochs,
            model_update=['stellar_model','ext_curve_b'],
            var_update = ['atm','E','plx'],
        )
        save_as_h5(ret, full_fn('hist_loss/hist_step0c.h5'))
        loss_hist.append(ret)

        # Plot stellar model
        print('Plotting stellar model ...')
        for i,track in enumerate(atm_tracks):
            fig,ax = model.plot_stellar_model(stellar_model, track)
            fig.savefig(full_fn(f'plots/stellar_model_step0c_track{i}'))
            plt.close(fig)

        fig,_ = plot_roughness(d_train['stellar_type'], stellar_model)
        fig.savefig(full_fn('plots/roughness_step0c'))
        plt.close(fig)

        fig,_ = model.plot_extinction_curve(stellar_model, show_variation=False)
        fig.savefig(full_fn('plots/extinction_curve_step0c'))
        plt.close(fig)

        np.save(full_fn('index/idx_good_wo_Rv.npy'), idx_good)
        stellar_model.save(full_fn('models/flux/xp_spectrum_model_final'))
        save_as_h5(d_train, full_fn('data/dtrain_final_wo_Rv.h5'))
        #save_as_h5(ret, full_fn('hist_loss/final_wo_Rv.h5'))

    if stage<2:

        d_val = load_h5(full_fn('data/d_val.h5'))
        d_train = load_h5(full_fn('data/dtrain_final_wo_Rv.h5'))
        atm_tracks = model.load_stellar_type_tracks(
            full_fn('models/prior/tracks.h5')
        )

        stellar_model = FluxModel.load(
            full_fn('models/flux/xp_spectrum_model_final-1')
        )

        print('Loading Gaussian Mixture Model prior on stellar type ...')
        stellar_type_prior = GaussianMixtureModel.load(
            full_fn('models/prior/gmm_prior-1')
        )

        # Initial weight of stars: equal
        weights_per_star = np.ones(len(d_train['plx']), dtype='f4')
        idx_hq = np.load(full_fn('index/idx_good_wo_Rv.npy'))

        # Run more steps in this stage
        n_epochs_1 = 256

        # Optimize the params of high-quality stars
        print('Optimizing params of hq stars ...')
        ret = train_stellar_model(
            stellar_model,
            d_train,
            weights_per_star,
            idx_train=np.where(idx_hq)[0],
            optimize_stellar_model=False,
            optimize_stellar_params=True,
            #lr_model_init=1e-7,
            #lr_stars_init=1e-5,
            batch_size=batch_size,
            n_epochs=n_epochs,
            var_update=['atm','E','plx'],
        )
        save_as_h5(ret, full_fn('hist_loss/hist_step1a.h5'))

        save_as_h5(d_train, full_fn('data/dtrain_final_wo_Rv_optimized.h5'))

        d_train = load_h5(full_fn('data/dtrain_final_wo_Rv_optimized.h5'))

        idx_large_E = (d_train['stellar_ext_est']>E_low)
        idx_hq_large_E = idx_hq & idx_large_E
        pct_use = np.mean(idx_hq_large_E)

        title = (
            r'$\mathrm{Training\ distribution'
            r'\ (step\ 1a:\ HQ,\ large\ E):'
            r'\ stellar\ type}$'
        )
        plot_param_histograms_1d(
            d_train['stellar_type'][idx_hq_large_E],
            weights_per_star[idx_hq_large_E],
            title,
            full_fn('plots/training_stellar_type_hist1d_step1a')
        )

        print(f'Learning (xi,E,plx) for {pct_use:.3%} of sources (high E) ...')
        ret = train_stellar_model(
            stellar_model,
            d_train,
            weights_per_star,
            idx_train=np.where(idx_hq_large_E)[0],
            optimize_stellar_model=False,
            optimize_stellar_params=True,
            #lr_model_init=1e-7,
            #lr_stars_init=1e-5,
            batch_size=batch_size,
            n_epochs=n_epochs,
            var_update=['E','plx','xi'],
        )

        save_as_h5(d_train, full_fn('data/dtrain_initial_guess_xi.h5'))

        d_train = load_h5(full_fn('data/dtrain_initial_guess_xi.h5'))

        weights_per_star = down_sample_weighing(
            d_train['xi_est'][idx_hq_large_E],
            d_train['xi_est'],
            bin_edges = np.linspace(-0.8, 0.8, n_bins+1)
        )
        weights_per_star *= weigh_prior(stellar_type_prior, d_train)
        weights_per_star = soft_clip_weights(weights_per_star, 10.)

        title = (
            r'$\mathrm{Training\ distribution'
            r'\ (step\ 1b:\ HQ,\ large\ E):'
            r'\ stellar\ type}$'
        )
        plot_param_histograms_1d(
            d_train['stellar_type'][idx_hq_large_E],
            weights_per_star[idx_hq_large_E],
            title,
            full_fn('plots/training_stellar_type_hist1d_step1b')
        )

        print(f'Learning (xi,E) for {pct_use:.3%} of sources (high E)')
        print('... in addition to mean extinction curve ...')
        ret = train_stellar_model(
            stellar_model,
            d_train,
            weights_per_star,
            idx_train=np.where(idx_hq_large_E)[0],
            optimize_stellar_model=True,
            optimize_stellar_params=True,
            lr_model_init=1e-5,
            lr_stars_init=1e-5,
            batch_size=batch_size,
            n_epochs=n_epochs,
            var_update=['E','xi'],
            model_update=['ext_curve_b'],
        )
        save_as_h5(ret, full_fn('hist_loss/hist_step1b.h5'))

        fig,_ = model.plot_extinction_curve(stellar_model, show_variation=True)
        fig.savefig(full_fn('plots/extinction_curve_step1b'))
        plt.close(fig)

        stellar_model.save(full_fn('models/flux/xp_spectrum_model_initial_update_b'))

        stellar_model = FluxModel.load(
            full_fn('models/flux/xp_spectrum_model_initial_update_b-1')
        )

        print(f'Learning (xi,E) for {pct_use:.3%} of sources (high E)')
        print('... in addition to extinction curve variation ...')
        ret = train_stellar_model(
            stellar_model,
            d_train,
            weights_per_star,
            idx_train=np.where(idx_hq_large_E)[0],
            optimize_stellar_model=True,
            optimize_stellar_params=True,
            #lr_model_init=5e-6,
            lr_stars_init=1e-5,
            batch_size=batch_size,
            n_epochs=n_epochs_1,
            var_update=['E','xi'],
            model_update=['ext_curve_w'],
        )
        save_as_h5(ret, full_fn('hist_loss/hist_step1c.h5'))

        fig,_ = model.plot_extinction_curve(stellar_model, show_variation=True)
        fig.savefig(full_fn('plots/extinction_curve_step1c'))
        plt.close(fig)

        np.save(full_fn('index/idx_hq.npy'), idx_hq)
        # Save initial guess of xi
        save_as_h5(d_train, full_fn('data/dtrain_Rv_initial.h5'))
        save_as_h5(ret, full_fn('hist_loss/Rv_initial.h5'))
        stellar_model.save(full_fn('models/flux/xp_spectrum_model_initial_Rv'))

    if stage<3:

        #n_epochs = 512

        stellar_model = FluxModel.load(
            full_fn('models/flux/xp_spectrum_model_initial_Rv-1')
        )
        d_train = load_h5(full_fn('data/dtrain_Rv_initial.h5'))
        d_val = load_h5(full_fn('data/d_val.h5'))

        idx_hq = np.load(full_fn('index/idx_hq.npy'))
        idx_large_E = (d_train['stellar_ext_est']>E_low)
        idx_hq_large_E = idx_hq & idx_large_E
        pct_use = np.mean(idx_hq_large_E)

        stellar_type_prior = GaussianMixtureModel.load(
            full_fn('models/prior/gmm_prior-1')
        )
        atm_tracks = model.calculate_stellar_type_tracks(stellar_type_prior)

        weights_per_star = down_sample_weighing(
            d_train['xi_est'][idx_hq_large_E],
            d_train['xi_est'],
            bin_edges = np.linspace(-0.8, 0.8, n_bins+1)
        )
        weights_per_star *= weigh_prior(stellar_type_prior, d_train)
        weights_per_star = soft_clip_weights(weights_per_star, 10.)

        print(f'Learning everything for {pct_use:.3%} of sources (high E) ...')
        ret = train_stellar_model(
            stellar_model,
            d_train,
            weights_per_star,
            idx_train=np.where(idx_hq_large_E)[0],
            optimize_stellar_model=True,
            optimize_stellar_params=True,
            lr_model_init=1e-5,
            lr_stars_init=1e-4,
            batch_size=batch_size,
            n_epochs=2*n_epochs,
            var_update=['atm','E','plx','xi'],
            model_update=['stellar_model','ext_curve_w','ext_curve_b'],
        )
        save_as_h5(ret, full_fn('hist_loss/hist_step2.h5'))

        print('Plotting stellar model ...')
        for i,track in enumerate(atm_tracks):
            fig,ax = model.plot_stellar_model(stellar_model, track)
            fig.savefig(full_fn(f'plots/stellar_model_step2_track{i}'))
            plt.close(fig)

        fig,_ = plot_roughness(d_train['stellar_type'], stellar_model)
        fig.savefig(full_fn('plots/roughness_step2'))
        plt.close(fig)

        fig,_ = model.plot_extinction_curve(stellar_model, show_variation=True)
        fig.savefig(full_fn('plots/extinction_curve_step2'))
        plt.close(fig)

        fig,_ = model.plot_RV_histogram(stellar_model, d_train)
        fig.savefig(full_fn('plots/RV_histogram_step2'))
        plt.close(fig)

        #fig,_ = model.plot_RV_skymap(stellar_model, d_train)
        #fig.savefig(full_fn('plots/RV_skymap_step2'))
        #plt.close(fig)

        save_as_h5(d_train, full_fn('data/dtrain_Rv_intermediate_0.h5'))
        save_as_h5(ret, full_fn('hist_loss/Rv_intermediate_0.h5'))
        stellar_model.save(
            full_fn('models/flux/xp_spectrum_model_intermediate_Rv')
        )

    if stage<4:

        stellar_model = FluxModel.load(
            full_fn('models/flux/xp_spectrum_model_intermediate_Rv-1')
        )
        stellar_type_prior = GaussianMixtureModel.load(
            full_fn('models/prior/gmm_prior-1')
        )

        d_train = load_h5(full_fn('data/dtrain_Rv_intermediate_0.h5'))
        d_val = load_h5(full_fn('data/d_val.h5'))
        atm_tracks = model.load_stellar_type_tracks(
            full_fn('models/prior/tracks.h5')
        )
        #idx_hq = np.load(full_fn('index/idx_hq.npy'))

        # Optimize all stellar params, in order to pick up
        # stars that were rejected due to extinction variation law

        weights_per_star = np.ones(d_train['stellar_type'].shape[0]).astype('float32')

        print('Learning stellar parameters for all sources ...')
        ret = train_stellar_model(
            stellar_model,
            d_train,
            weights_per_star,
            #idx_train=idx_hq_large_E,
            optimize_stellar_model=False,
            optimize_stellar_params=True,
            batch_size=batch_size,
            n_epochs=n_epochs//2,
            var_update = ['atm','E','plx','xi'],
        )

        save_as_h5(d_train, full_fn('data/dtrain_Rv_intermediate_1.h5'))
        save_as_h5(ret, full_fn('hist_loss/Rv_intermediate_1.h5'))

        d_train = load_h5(full_fn('data/dtrain_Rv_intermediate_1.h5'))
        d_val = load_h5(full_fn('data/d_val.h5'))

        # Remove outliers
        idx_params_good = identify_outlier_stars( # TODO: rename to "identify_param_outliers"
            d_train,
            sigma_clip_teff=4., # TODO: Make this clipping be agnostic about the parameter names
            sigma_clip_logg=4.,
            sigma_clip_feh=4.,
        )
        pct_good = np.mean(idx_params_good)
        print(f'Parameter outliers: {1-pct_good:.3%} of sources.')

        idx_flux_good = identify_flux_outliers(
            d_train, stellar_model,
            chi2_dof_clip=10.,
            #chi_indiv_clip=20.
        )
        pct_good = np.mean(idx_flux_good)
        print(f'Flux outliers: {1-pct_good:.3%} of sources.')

        idx_good = idx_params_good & idx_flux_good
        pct_good = np.mean(idx_good)
        print(f'Combined outliers: {1-pct_good:.3%} of sources.')

        idx_final_train = idx_good #& (d_train['stellar_ext_est']>0.1)
        pct_use = np.mean(idx_final_train)
        print(f'Training on {pct_use:.3%} of sources.')

        weights_per_star = down_sample_weighing(
            d_train['xi_est'][idx_hq_large_E],
            d_train['xi_est'],
            bin_edges = np.linspace(-0.8, 0.8, n_bins+1)
        )
        weights_per_star *= weigh_prior(stellar_type_prior, d_train)
        weights_per_star = soft_clip_weights(weights_per_star, 10.)

        title = (
            r'$\mathrm{Training\ distribution'
            r'\ (step\ 3a:\ cut\ flux/param\ outliers):'
            r'\ stellar\ type}$'
        )
        plot_param_histograms_1d(
            d_train['stellar_type'][idx_final_train],
            weights_per_star[idx_final_train],
            title,
            full_fn('plots/training_stellar_type_hist1d_step3a')
        )

        pct_use = np.mean(idx_final_train)
        print('Learning everything except extinction curve variation')
        print(f'using {pct_use:.3%} of sources (cut flux/param outliers) ...')
        ret = train_stellar_model(
            stellar_model,
            d_train,
            np.ones_like(weights_per_star, dtype='f4'),
            idx_train=np.where(idx_final_train)[0],
            optimize_stellar_model=True,
            optimize_stellar_params=True,
            lr_model_init=1e-5,
            lr_stars_init=1e-4,
            batch_size=batch_size,
            n_epochs=n_epochs//2,
            # Freeze extinction curve variation when training with
            # all stars (not necessarily having high E)
            var_update=['atm','E','plx','xi'],
            model_update=['stellar_model','ext_curve_b'],
        )

        stellar_model.save(full_fn('models/flux/xp_spectrum_model_final_Rv'))
        save_as_h5(ret, full_fn('hist_loss/final_Rv.h5'))

        fig,ax = model.plot_stellar_model(stellar_model, track)
        fig.savefig(full_fn(f'plots/stellar_model_step3_track{i}'))
        plt.close(fig)

        fig,_ = plot_roughness(d_train['stellar_type'], stellar_model)
        fig.savefig(full_fn('plots/roughness_step3'))
        plt.close(fig)

        fig,_ = model.plot_extinction_curve(stellar_model, show_variation=True)
        fig.savefig(full_fn('plots/extinction_curve_step3'))
        plt.close(fig)

        fig,_ = model.plot_RV_histogram(stellar_model, d_train)
        fig.savefig(full_fn('plots/RV_histogram_step3'))
        plt.close(fig)

        fig,_ = model.plot_RV_skymap(stellar_model, d_train)
        fig.savefig(full_fn('plots/RV_skymap_step3'))
        plt.close(fig)

        stellar_model.save(full_fn('models/flux/xp_spectrum_model_final_Rv'))
        save_as_h5(ret, full_fn('hist_loss/final_Rv.h5'))

        print('Optimizing stellar parameters in training set ...')
        ret = train_stellar_model(
            stellar_model,
            d_train,
            weights_per_star,
            #idx_train=np.where(idx_final_train)[0],
            optimize_stellar_model=False,
            optimize_stellar_params=True,
            batch_size=batch_size,
            n_epochs=n_epochs,
            var_update = ['atm','E','plx','xi'],
        )

        fig,_ = model.plot_RV_histogram(stellar_model, d_train)
        fig.savefig(full_fn('plots/RV_histogram_final'))
        plt.close(fig)

        fig,_ = model.plot_RV_skymap(stellar_model, d_train)
        fig.savefig(full_fn('plots/RV_skymap_final'))
        plt.close(fig)

        for key in ('stellar_type', 'xi', 'stellar_ext', 'plx'):
            d_val[f'{key}_est'] = d_val[key].copy()

        print('Optimizing stellar parameters in validation set ...')
        ret = train_stellar_model(
            stellar_model,
            d_val,
            weights_per_star,
            #idx_train=np.where(idx_final_train)[0],
            optimize_stellar_model=False,
            optimize_stellar_params=True,
            batch_size=batch_size,
            n_epochs=n_epochs,
            var_update = ['atm','E','plx','xi'],
        )

        save_as_h5(d_train, full_fn('data/dtrain_Rv_final.h5'))
        save_as_h5(d_val, full_fn('data/dval_Rv_final.h5'))

        return 0

def plot_roughness(theta, stellar_model):
    roughness_plot = batch_apply_tf(
        lambda x: stellar_model.roughness(tf.convert_to_tensor(x),
                                          reduce_stars=False),
        1024, theta,
        function=True,
        numpy=True,
        progress=False
    )
    fig,ax_arr = plot_utils.projection_grid(
        theta,
        c=roughness_plot,
        labels=(r'$T_{\mathrm{eff}} / (1000\,\mathrm{K})$',
                r'$\left[\mathrm{Fe/H}\right] / \mathrm{dex}$',
                r'$\log g / \mathrm{dex}$'),
        extents=[(13.,3.),(-2.7,1.0),(5.2,-0.3)],
        scatter=False,
        clabel=(r'$\langle\left| '
                r'\nabla_{\!\theta}\, f_{\lambda}\, '
                r'\right|^2\rangle$'),
        hist_kw=dict(bins=100),
        #fig_kw=dict(dpi=200),
        norm=Normalize(0,5)
    )
    for ax in ax_arr.flat:
        ax.xaxis.set_minor_locator(AutoMinorLocator())
        ax.yaxis.set_minor_locator(AutoMinorLocator())
    rough = np.std(roughness_plot)
    fig.suptitle(rf'$\mathrm{{roughness}} = {rough:.4f}$')
    fig.subplots_adjust(
        left=0.14, right=0.97,
        bottom=0.12, top=0.92,
        hspace=0.07, wspace=0.07
    )
    return fig,ax_arr

def plot_camd(d, color_key=None, color_dim=None, low_extinction=False):
    fig,ax = plt.subplots(figsize=(6,6), layout='constrained')

    # Calculate M_G and BP-RP
    zp_g,zp_bp,zp_rp = (25.6884, 25.3514, 24.7619)
    m_g = -2.5*np.log10(d['phot_g_mean_flux']) + zp_g
    m_bp = -2.5*np.log10(d['phot_bp_mean_flux']) + zp_bp
    m_rp = -2.5*np.log10(d['phot_rp_mean_flux']) + zp_rp
    bprp = m_bp - m_rp
    g_absmag = m_g + 5*np.log10(d['plx']) - 10.

    # Only plot stars with well determined Gaia parallaxes
    idx = (d['plx'] / d['plx_err'] > 3.)

    # Additionally require decent photometry
    for b in ('g','bp','rp'):
        idx &= (d[f'phot_{b}_mean_flux']/d[f'phot_{b}_mean_flux_error']>5)

    # Select low-extinction stars?
    if low_extinction:
        idx &= (d['stellar_ext'] + d['stellar_ext_err'] < 0.3)

    bprp = bprp[idx]
    g_absmag = g_absmag[idx]

    bprp_lim = (-1.0, 5.0)
    g_absmag_lim = (15., -5.5)

    if color_key is None:
        _,_,_,im = ax.hist2d(
            bprp, g_absmag,
            range=[sorted(bprp_lim),sorted(g_absmag_lim)],
            bins=50,
            norm=LogNorm()
        )
        ax.set_xlim(bprp_lim)
        ax.set_ylim(g_absmag_lim)
        clabel = r'$\mathrm{density}$'
    else:
        c = d[color_key][idx]

        if color_dim is not None:
            c = c[:,color_dim]
        clim = plot_utils.choose_lim(c)

        im = plot_utils.hist2d_reduce(
            bprp, g_absmag, c,
            ax=ax,
            xlim=bprp_lim,
            ylim=g_absmag_lim,
            hist_kw=dict(bins=50),
            imshow_kw=dict(vmin=clim[0],vmax=clim[1])
        )

        clabel = color_key.replace(r'_', r'\_')
        if color_dim is not None:
            clabel = rf'{clabel}\_{color_dim}'
        clabel = rf'$\mathtt{{{clabel}}}$'

    cb = fig.colorbar(im, ax=ax, label=clabel)
    ax.set_xlabel(r'$BP-RP$')
    ax.set_ylabel(r'$m_G + 5\log_{10}\hat{\varpi} - 10$')

    return fig,ax

def plot_training_data(d, output_dir=''):
    fig,ax = plot_camd(d)
    fn = os.path.join(output_dir, 'plots', f'camd_density')
    fig.savefig(fn)
    plt.close(fig)

    n_params = d['stellar_type'].shape[1]
    for key in ('stellar_type','stellar_type_err'):
        for dim in range(n_params):
            print(f'Plotting {key}_{dim} ...')
            fig,ax = plot_camd(
                d, color_key=key, color_dim=dim,
                low_extinction=True
            )
            fn = os.path.join(output_dir, 'plots', f'camd_{key}_{dim}')
            fig.savefig(fn)
            plt.close(fig)

    print(f'Plotting extinctions ...')
    fig,ax = plot_camd(d, color_key='stellar_ext')
    fn = os.path.join(output_dir, 'plots', 'camd_extinction')
    fig.savefig(fn)
    plt.close(fig)

def main():
    parser = ArgumentParser(
        description='Train forward model of XP spectra.',
        add_help=True
    )
    parser.add_argument(
        '-i', '--input',
        type=str,
        required=True,
        help='Filename of training data.'
    )
    parser.add_argument(
        '-o', '--output-dir',
        type=str,
        required=True,
        help='Directory in which to store output.'
    )
    parser.add_argument(
        '--stage',
        type=int,
        default=0,
        help='Stage at which to begin training.'
    )
    parser.add_argument(
        '--thin',
        type=int,
        default=1,
        help='Only use every Nth star in the training set.'
    )
    parser.add_argument(
        '--E_low',
        type=float,
        default=0.1,
        help='Lower limit of extinction for ext curve variation.'
    )

    args = parser.parse_args()
    train(args.input, args.output_dir, stage=args.stage, thin=args.thin, E_low=args.E_low)

    return 0

if __name__=='__main__':
    main()


#!/usr/bin/env python
import os
import argparse

import numpy as np
import cv2 as cv

import patched_cnn as pcnn
import image_processor as ip

import copy

parser = argparse.ArgumentParser()
parser.add_argument("--image", help="path to image.")
parser.add_argument("--model", help="path to trained model h5 file.")
args = parser.parse_args()

image_source = args.image
model_source = args.model

img = cv.imread(image_source)

imgLab = cv.cvtColor(img, cv.COLOR_BGR2LAB)

# print(np.concatenate(( img, imgLab ), 2).shape)

input = np.array([ cv.resize( imgLab, (32, 32) ) ]) # input to nn should be (1, 32, 32, 3)

cnn = pcnn.Patched_CNN()
cnn.load_model(model_source)
print("...initialized model...")

pred = cnn.predict(input)

shadow_mask = cv.resize( pred.reshape(32, 32), (img.shape[1], img.shape[0]) )

outImg = copy.deepcopy(img)

imgp = ip.ImageProcessed(img, "image_name")
imgp.shadow_mask = shadow_mask
imgp._segment()
imgp.label_shadow_segments(0.1)
imgp.showShadows()

for col in range(outImg.shape[0]):
    for row in range(outImg.shape[1]):
        # if shadow_mask[col, row] > 0.5:
            # outImg[col, row][0] = 255.0 * shadow_mask[col, row]
            # outImg[col, row][0] = outImg[col, row][0] * (1-shadow_mask[col, row])
            outImg[col, row][1] = 255 * (shadow_mask[col, row])
            # outImg[col, row][2] = 255.0 * shadow_mask[col, row]

# cv.namedWindow("original", cv.WINDOW_NORMAL)
cv.namedWindow("mask", cv.WINDOW_NORMAL)

# cv.imshow("original", img)
cv.imshow("mask", outImg)
# # cv.imshow("predicted shadow mask", shadow_mask)
cv.waitKey()

# for i in range(len(img_processed.segments)):
#     img_processed.segments[i]['is_shadow'] = pred[i] == 1

# img_processed.showShadows()


import cv2
from PIL import Image

from utils import get_limits

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

    image = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)

    cv2.imshow('frame', frame)

    if cv2.waitKey(1) & 0xFF == ord('q'):  # Continuous display; press 'q' to exit
        break

cap.release()

cv2.destroyAllWindows()


from easydict import EasyDict

configer = EasyDict()

configer.dsize = (64, 64)
configer.n_channel = 23
configer.n_class = 63

configer.splitmode = 'split_{}x{}_1'.format(configer.dsize[0], configer.dsize[1])
configer.modelbase = 'recognize_shufflev2'

configer.datatype = 'Multi'
if configer.datatype == 'Multi':
    configer.usedChannels = [550]
    configer.n_usedChannels = len(configer.usedChannels)
    configer.modelname = '{}_{}_{}chs_{}sta_20nm'.\
                    format(configer.modelbase, configer.splitmode, configer.n_usedChannels, configer.usedChannels[0])
elif configer.datatype == 'RGB':
    configer.usedChannels = 'R'
    configer.n_usedChannels = len(configer.usedChannels)
    configer.modelname = '{}_{}_{}'.\
                    format(configer.modelbase, configer.splitmode, configer.usedChannels)

configer.datapath = '/datasets/ECUST2019_{}x{}'.\
                                format(configer.dsize[0], configer.dsize[1])
configer.logspath = './logspath'
configer.mdlspath = './mdlpath'

## training step
configer.batchsize = 64
configer.n_epoch   = 500

## learing rate
configer.lrbase = 0.001
configer.stepsize = 300
configer.gamma = 0.1

configer.cuda = True


import itertools
from collections.abc import Iterable

import numpy as np

from ... import opcodes as OperandDef
from ...serialization.serializables import TupleField
from ...config import options
from ..utils import decide_chunk_sizes, gen_random_seeds
from .core import TensorRandomOperandMixin, TensorDistribution

class TensorDirichlet(TensorDistribution, TensorRandomOperandMixin):
    _op_type_ = OperandDef.RAND_DIRICHLET

    _fields_ = "alpha", "size"
    alpha = TupleField("alpha", default=None)
    _func_name = "dirichlet"

    def _calc_shape(self, shapes):
        shape = super()._calc_shape(shapes)
        return shape + (len(self.alpha),)

    def __call__(self, chunk_size=None):
        return self.new_tensor(None, None, raw_chunk_size=chunk_size)

    @classmethod
    def tile(cls, op):
        tensor = op.outputs[0]
        chunk_size = tensor.extra_params.raw_chunk_size or options.chunk_size
        nsplits = decide_chunk_sizes(
            tensor.shape[:-1], chunk_size, tensor.dtype.itemsize
        )
        nsplits += ((len(op.alpha),),)

        idxes = list(itertools.product(*[range(len(s)) for s in nsplits]))
        seeds = gen_random_seeds(len(idxes), np.random.RandomState(op.seed))

        out_chunks = []
        for seed, idx, shape in zip(seeds, idxes, itertools.product(*nsplits)):
            inputs = [inp.cix[idx] for inp in op.inputs]
            size = shape[:-1]

            chunk_op = op.copy().reset_key()
            chunk_op._state = None
            chunk_op.seed = seed
            chunk_op.size = size
            out_chunk = chunk_op.new_chunk(inputs, shape=shape, index=idx)
            out_chunks.append(out_chunk)

        new_op = op.copy()
        return new_op.new_tensors(
            op.inputs, tensor.shape, chunks=out_chunks, nsplits=nsplits
        )

def dirichlet(random_state, alpha, size=None, chunk_size=None, gpu=None, dtype=None):
    r"""
    Draw samples from the Dirichlet distribution.

    Draw `size` samples of dimension k from a Dirichlet distribution. A
    Dirichlet-distributed random variable can be seen as a multivariate
    generalization of a Beta distribution. Dirichlet pdf is the conjugate
    prior of a multinomial in Bayesian inference.

    Parameters
    ----------
    alpha : array
        Parameter of the distribution (k dimension for sample of
        dimension k).
    size : int or tuple of ints, optional
        Output shape.  If the given shape is, e.g., ``(m, n, k)``, then
        ``m * n * k`` samples are drawn.  Default is None, in which case a
        single value is returned.
    chunk_size : int or tuple of int or tuple of ints, optional
        Desired chunk size on each dimension
    gpu : bool, optional
        Allocate the tensor on GPU if True, False as default
    dtype : data-type, optional
      Data-type of the returned tensor.

    Returns
    -------
    samples : Tensor
        The drawn samples, of shape (size, alpha.ndim).

    Raises
    -------
    ValueError
        If any value in alpha is less than or equal to zero

    Notes
    -----
    .. math:: X \approx \prod_{i=1}^{k}{x^{\alpha_i-1}_i}

    Uses the following property for computation: for each dimension,
    draw a random sample y_i from a standard gamma generator of shape
    `alpha_i`, then
    :math:`X = \frac{1}{\sum_{i=1}^k{y_i}} (y_1, \ldots, y_n)` is
    Dirichlet distributed.

    References
    ----------
    .. [1] David McKay, "Information Theory, Inference and Learning
           Algorithms," chapter 23,
           http://www.inference.phy.cam.ac.uk/mackay/
    .. [2] Wikipedia, "Dirichlet distribution",
           http://en.wikipedia.org/wiki/Dirichlet_distribution

    Examples
    --------
    Taking an example cited in Wikipedia, this distribution can be used if
    one wanted to cut strings (each of initial length 1.0) into K pieces
    with different lengths, where each piece had, on average, a designated
    average length, but allowing some variation in the relative sizes of
    the pieces.

    >>> import mars.tensor as mt

    >>> s = mt.random.dirichlet((10, 5, 3), 20).transpose()

    >>> import matplotlib.pyplot as plt

    >>> plt.barh(range(20), s[0].execute())
    >>> plt.barh(range(20), s[1].execute(), left=s[0].execute(), color='g')
    >>> plt.barh(range(20), s[2].execute(), left=(s[0]+s[1]).execute(), color='r')
    >>> plt.title("Lengths of Strings")
    """
    if isinstance(alpha, Iterable):
        alpha = tuple(alpha)
    else:
        raise TypeError("`alpha` should be an array")
    if dtype is None:
        dtype = np.random.RandomState().dirichlet(alpha, size=(0,)).dtype
    size = random_state._handle_size(size)
    seed = gen_random_seeds(1, random_state.to_numpy())[0]
    op = TensorDirichlet(seed=seed, alpha=alpha, size=size, gpu=gpu, dtype=dtype)
    return op(chunk_size=chunk_size)


# https://www.acmicpc.net/problem/1987

import sys
from collections import deque

def bfs(x: int, y: int) -> int:
    answer = 1
    queue = deque([(x, y, arr[x][y])])

    while queue:
        qx, qy, alphas = queue.pop()

        for i in range(4):
            nx = qx + dx[i]
            ny = qy + dy[i]

            if 0 <= nx < r and 0 <= ny < c and not arr[nx][ny] in alphas:
                queue.append((nx, ny, alphas + arr[nx][ny]))
                answer = max(answer, len(alphas) + 1)

    return answer

r, c = map(int, sys.stdin.readline().split())
arr = [list(sys.stdin.readline().rstrip()) for _ in range(r)]
dx = [-1, 1, 0, 0]
dy = [0, 0, -1, 1]

print(bfs(0, 0))


from django.shortcuts import redirect, render

from ..forms import ArtistForm

__all__ = (
    'artist_add',
)

def artist_add(request):
    """
    1. artist_add.html 작성
    2. url과 연결, /artist/add/에 매핑
    3. GET요청 시 잘 되는지 확인
    4. form method설정  후 POST요청시 artist_add() view 분기
    5. POST 요청의 값이 request.POST에 잘 오는지 확인
        name값만 받아서 name만 갖는 Artist를 먼저 생성
        성공시 나머지 값들을 하나씩 적용해보기
    6. request.POST에 담긴 값을 사용해 Artist인스턴스 생성
    7. 생성 완료후 'artist:artist-list' URL에 해당하는 view로 이동
    """
    # if request.method == "POST":
    #     name = request.POST['name']
    #     Artist.objects.create(
    #         name=name,
    #     )
    #     return redirect('artist:artist-list')
    # else:
    #     return render(request, 'artist/artist_add.html')
    if request.method == "POST":
        form = ArtistForm(request.POST, request.FILES)
        form.save()
        return redirect('artist:artist-list')
    else:
        form = ArtistForm()
    context = {
        'artist_form': form,
    }
    return render(request, 'artist/artist_add.html', context)


from django.urls import include, path
from . import views
from . import api

from rest_framework import routers

router = routers.DefaultRouter()
router.register(r'pools', api.PoolsViewSet, basename='pools')
router.register(r'questions', api.QuestionViewSet, basename='questions')

app_name = 'poolsapp'

urlpatterns = [
    path('', views.index, name='index'),
    path('create/', views.create_pool, name='create_pool'),
    path('pool/<int:pool_id>', views.pool ,name='pool'),

    #API
    path('api/', include(router.urls)),
]


from io import BytesIO
import os
from django.db import models
import qrcode
from dqr.settings import MEDIA_URL
from django.core.files import File

# Create your models here.
class Qrcode(models.Model):
    name = models.CharField(max_length=200)
    internal_url = models.CharField(max_length=500)
    external_url = models.CharField(max_length=500)
    modified = models.DateTimeField(auto_now=True)
    image = models.ImageField(upload_to='qrcodes/', blank=True)

    def make_qrcode(self):
        qr = qrcode.QRCode(
            version=1,
            error_correction=qrcode.constants.ERROR_CORRECT_L,
            box_size=10,
            border=4,
        )
        qr.add_data("http://dqr.michaelbasta.it/redirect/" + self.internal_url)
        qr.make(fit=True)
        qr_image = qr.make_image(fill_color="black", back_color="transparent")
        image_url = os.path.join(MEDIA_URL, 'qrcodes/', self.internal_url + '.png')

        # Create a BytesIO buffer to hold the image data
        buffer = BytesIO()
        qr_image.save(buffer)
        buffer.seek(0)

        # Save the image to the database
        self.image.save(self.internal_url + '.png', File(buffer))
        print(self.image.url)
        print(self.image.path)
        self.save()

    def __str__(self):
        return self.id +' - ' + self.name + ' - ' + self.internal_url + ' - ' + self.external_url + ' - ' + self.modified


from PIL import Image

gif = []
files = []
for n in range(0,128):
    files.append(f"dump/cot/{n}.png")

# using https://stackoverflow.com/a/451580/17091581
# rescale
rescale = False
if rescale:
    basewidth = 1200
    for f in files:
        picture = Image.open(f)
        wpercent = (basewidth/float(picture.size[0]))
        hsize = int((float(picture.size[1])*float(wpercent)))
        picture = picture.resize((basewidth,hsize), Image.Resampling.LANCZOS)
        gif.append(picture)
else:
    for f in files:
        gif.append(Image.open(f))

gif[0].save('betterztoc.gif', save_all=True, optimize=False, append_images=gif[1:], loop=0, )


import matplotlib.pyplot as plt
import numpy as np
from funcsforprajay.plotting.plotting import plot_bar_with_points

import _utils_.alloptical_plotting as aoplot

# import onePstim superobject that will collect analyses from various individual experiments
from _exp_metainfo_.exp_metainfo import import_resultsobj, ExpMetainfo
from _utils_.io import import_expobj
from onePexperiment.OnePhotonStimAnalysis_main import OnePhotonStimAnalysisFuncs
from onePexperiment.OnePhotonStimMain import OnePhotonStim

results_object_path = '/home/pshah/mnt/qnap/Analysis/onePstim_results_superobject.pkl'
onePresults = import_resultsobj(pkl_path=results_object_path)

date = '2021-01-24'

expobj = import_expobj(prep='PS11', trial='t-012', date=date)  # post4ap trial

# %% B) LFP signal with optogenetic stims

date = '2021-01-24'

# pre4ap
expobj: OnePhotonStim = import_expobj(prep='PS11', trial='t-009', date=date)  # pre4ap trial
aoplot.plotLfpSignal(expobj, x_axis='time', figsize=(5 / 100 * 100, 3), linewidth=0.5, downsample=True,
                     sz_markings=False, color='black',
                     ylims=[-4, 1], xlims=[105 * expobj.paq_rate, 205 * expobj.paq_rate])

# post4ap
expobj: OnePhotonStim = import_expobj(prep='PS11', trial='t-012', date=date)  # post4ap trial
aoplot.plotLfpSignal(expobj, x_axis='time', figsize=(5 / 100 * 150, 3), linewidth=0.5, downsample=True,
                     sz_markings=False, color='black',
                     ylims=[0, 5], xlims=[10 * expobj.paq_rate, 160 * expobj.paq_rate])

# Avg Flu signal with optogenetic stims

offset = expobj.frame_start_time_actual / expobj.paq_rate

# pre4ap
expobj = import_expobj(prep='PS11', trial='t-009', date=date)  # pre4ap trial
aoplot.plotMeanRawFluTrace(expobj, stim_span_color='cornflowerblue', x_axis='Time (secs)',
                           xlims=[105 * expobj.fps, 205 * expobj.fps], stim_lines=False,
                           figsize=[5 / 100 * 100, 3])

# post4ap
expobj = import_expobj(prep='PS11', trial='t-012', date=date)  # post4ap trial
aoplot.plotMeanRawFluTrace(expobj, stim_span_color='cornflowerblue', x_axis='Time (secs)',
                           xlims=[10 * expobj.fps, 160 * expobj.fps], stim_lines=False,
                           figsize=[5 / 100 * 150, 3])

# %% C) avg LFP trace 1p stim plots

date = '2021-01-24'

# pre4ap
pre4ap = import_expobj(prep='PS11', trial='t-009', date=date)  # pre4ap trial

assert 'pre' in pre4ap.exptype
aoplot.plot_flu_1pstim_avg_trace(pre4ap, x_axis='time', individual_traces=False, stim_span_color='skyblue',
                                 y_axis='dff', quantify=False, figsize=[3, 3], title='Baseline', ylims=[-0.5, 2.0])

aoplot.plot_lfp_1pstim_avg_trace(pre4ap, x_axis='time', individual_traces=False, pre_stim=0.25, post_stim=0.75,
                                 write_full_text=False,
                                 optoloopback=True, figsize=(3.1, 3), shrink_text=0.8,
                                 stims_to_analyze=pre4ap.stim_start_frames,
                                 title='Baseline')

# post4ap
post4ap = import_expobj(prep='PS11', trial='t-012', date=date)  # post4ap trial

assert 'post' in post4ap.exptype
aoplot.plot_flu_1pstim_avg_trace(post4ap, x_axis='time', individual_traces=False, stim_span_color='skyblue',
                                 stims_to_analyze=post4ap.stims_out_sz,
                                 y_axis='dff', quantify=False, figsize=[3, 3], title='Interictal', ylims=[-0.5, 2.0])

aoplot.plot_lfp_1pstim_avg_trace(post4ap, x_axis='time', individual_traces=False, pre_stim=0.25, post_stim=0.75,
                                 write_full_text=False,
                                 optoloopback=True, figsize=(3.1, 3), shrink_text=0.8,
                                 stims_to_analyze=post4ap.stims_out_sz,
                                 title='Interictal')

aoplot.plot_flu_1pstim_avg_trace(post4ap, x_axis='time', individual_traces=False, stim_span_color='skyblue',
                                 stims_to_analyze=post4ap.stims_in_sz,
                                 y_axis='dff', quantify=False, figsize=[3, 3], title='Ictal', ylims=[-0.5, 2.0])

aoplot.plot_lfp_1pstim_avg_trace(post4ap, x_axis='time', individual_traces=False, pre_stim=0.25, post_stim=0.75,
                                 write_full_text=False,
                                 optoloopback=True, figsize=(3.1, 3), shrink_text=0.8,
                                 stims_to_analyze=post4ap.stims_in_sz,
                                 title='Ictal')

# %% D) BAR PLOT OF RESPONSE MAGNITUDE FOR 1P STIM EXPERIMENTS

data = [[rp for rp in onePresults.mean_stim_responses.iloc[:, 1] if rp != '-']]
data.append([rp for rp in onePresults.mean_stim_responses.iloc[:, 2] if rp != '-'])
# data.append([rp for rp in onePresults.mean_stim_responses.iloc[:,3] if rp != '-'])

interictal_response_magnitude = {
    'PS07': [0.2654, 0.1296],
    'PS11': [0.9005, 0.9208],
    'PS18': [1.5309],
    'PS09': [0.2543, 0.6789, 0.763],
    'PS16': [0.1858, 0.2652]
}

baseline_response_magnitude = {
    # 'PS17': [0.1686],  # can't be used for paird pre4ap to interictal plot
    'PS07': [0.3574, 0.3032],
    'PS11': [0.4180, 0.3852],
    'PS18': [0.3594],
    'PS09': [0.2662],
    'PS16': [0.1217, 0.093, 0.087]
}

interictal_response_magnitude_plot = [np.mean(items) for items in interictal_response_magnitude.values()]
baseline_response_magnitude_plot = [np.mean(items) for items in baseline_response_magnitude.values()]

fig, ax = plt.subplots(figsize=[3, 5], dpi = 300)
plot_bar_with_points(data=[baseline_response_magnitude_plot, interictal_response_magnitude_plot],
                     title='response magnitudes', x_tick_labels=['baseline', 'interictal'], paired=True,
                     points=True, bar=False, colors=['gray', 'green'], fig=fig, ax=ax, show=False, s=50,
                     x_label='experiment groups', y_label='Avg. dFF', alpha=1,
                     expand_size_x=0.5, expand_size_y=1.3, shrink_text=1.35, ylims=[0, 2.5])
fig.tight_layout(pad=0.2)
fig.show()
save_path = '/home/pshah/mnt/qnap/Analysis/' + 'onePstim_response_quant'
print('saving fig to: ', save_path)
fig.savefig(fname=save_path + '.png', transparent=True, format='png')
fig.savefig(fname=save_path + '.svg', transparent=True, format='svg')

# todo run paired ttest stats

# %% E) BAR PLOT OF DECAY CONSTANT FOR 1P STIM EXPERIMENTS

data = [list(onePresults.mean_stim_responses[onePresults.mean_stim_responses.iloc[:, -3].notnull()].iloc[:, -3])]
data.append(list(onePresults.mean_stim_responses[onePresults.mean_stim_responses.iloc[:, -2].notnull()].iloc[:, -2]))
# data.append(ls(onePresults.mean_stim_responses[onePresults.mean_stim_responses.iloc[:, -1].notnull()].iloc[:, -1]))

interictal_decay_constant = {
    'PS07': [1.26471, 0.32844],
    'PS11': [1.05109, 0.45986],
    'PS18': [0.65681],
    'PS09': [0.52549, 0.59118, 0.39412],
    'PS16': [0.78811, 0.98517]
}

baseline_decay_constant = {
    # 'PS17': [0.1686],
    'PS07': [0.45972, 0.46595],
    'PS11': [0.45980, 0.52553],
    'PS18': [0.45975],
    'PS09': [0.52546],
    'PS16': [0.45964, 0.39399, 0.46595]
}

baseline_decay_constant_plot = [np.mean(items) for items in baseline_decay_constant.values()]
interictal_decay_constant_plot = [np.mean(items) for items in interictal_decay_constant.values()]

fig, ax = plt.subplots(figsize=[3, 5])
plot_bar_with_points(data=[baseline_decay_constant_plot, interictal_decay_constant_plot], title='decay constants',
                     legend_labels=list(onePresults.mean_stim_responses.columns[-3:]), paired=True, x_tick_labels=['baseline', 'interictal'],
                     points=True, bar=False, colors=['gray', 'green'], fig=fig, ax=ax, show=False,
                     x_label='experiment groups', y_label='Avg. Decay constant (secs.)', alpha=1, s=50,
                     expand_size_x=0.9, expand_size_y=1.2, shrink_text=1.35, ylims=[0.2, 1.1])
fig.tight_layout(pad=0.2)
fig.show()
save_path = '/home/pshah/mnt/qnap/Analysis/' + 'onePstim_decay_quant'
print('saving fig to: ', save_path)
fig.savefig(fname=save_path + '.png', transparent=True, format='png')
fig.savefig(fname=save_path + '.svg', transparent=True, format='svg')

# todo run paired ttest stats

# %% F) Radial plot of Mean FOV for photostimulation trials, with period equal to that of photostimulation timing period

# run data analysis
exp_sz_occurrence = OnePhotonStimAnalysisFuncs.collectSzOccurrenceRelativeStim()

# make plot
bin_width = int(1 * expobj.fps)
period = len(np.arange(0, (expobj.stim_interval_fr / bin_width)))
theta = (2 * np.pi) * np.arange(0, (expobj.stim_interval_fr / bin_width)) / period

fig, ax = plt.subplots(subplot_kw={'projection': 'polar'}, dpi=600)
for exp, values in exp_sz_occurrence.items():
    plot = values
    ax.bar(theta, plot, width=(2 * np.pi) / period, bottom=0.0, alpha=0.5)
ax.set_rmax(1.1)
# ax.set_rticks([1])  # Less radial ticks
ax.set_rlabel_position(-35.5)  # Move radial labels away from plotted line
ax.grid(True)
ax.set_xticks((2 * np.pi) * np.arange(0, (expobj.stim_interval_fr / bin_width)) / period)
ax.set_title("sz probability occurrence (binned every 1s)", va='bottom')
ax.spines['polar'].set_visible(False)
fig.show()


import keras.backend as K
import numpy as np
from keras.engine.topology import InputSpec, Layer
from keras.layers import (Activation, Concatenate, Conv2D, Flatten, Input,
                          Reshape)
from keras.models import Model

from nets.vgg import VGG16

class Normalize(Layer):
    def __init__(self, scale, **kwargs):
        self.axis = 3
        self.scale = scale
        super(Normalize, self).__init__(**kwargs)

    def build(self, input_shape):
        self.input_spec = [InputSpec(shape=input_shape)]
        shape = (input_shape[self.axis],)
        init_gamma = self.scale * np.ones(shape)
        self.gamma = K.variable(init_gamma, name='{}_gamma'.format(self.name))
        self.trainable_weights = [self.gamma]

    def call(self, x, mask=None):
        output = K.l2_normalize(x, self.axis)
        output *= self.gamma
        return output

def SSD300(input_shape, num_classes=21):
    #---------------------------------#
    #   典型的输入大小为[300,300,3]
    #---------------------------------#
    input_tensor = Input(shape=input_shape)

    #------------------------------------------------------------------------#
    #   net变量里面包含了整个SSD的结构，通过层名可以找到对应的特征层
    #------------------------------------------------------------------------#
    net = VGG16(input_tensor)

    #-----------------------将提取到的主干特征进行处理---------------------------#
    # 对conv4_3的通道进行l2标准化处理
    # 38,38,512
    net['conv4_3_norm'] = Normalize(20, name='conv4_3_norm')(net['conv4_3'])
    num_anchors = 4
    # 预测框的处理
    # num_anchors表示每个网格点先验框的数量，4是x,y,h,w的调整
    net['conv4_3_norm_mbox_loc']        = Conv2D(num_anchors * 4, kernel_size=(3,3), padding='same', name='conv4_3_norm_mbox_loc')(net['conv4_3_norm'])
    net['conv4_3_norm_mbox_loc_flat']   = Flatten(name='conv4_3_norm_mbox_loc_flat')(net['conv4_3_norm_mbox_loc'])
    # num_anchors表示每个网格点先验框的数量，num_classes是所分的类
    net['conv4_3_norm_mbox_conf']       = Conv2D(num_anchors * num_classes, kernel_size=(3,3), padding='same',name='conv4_3_norm_mbox_conf')(net['conv4_3_norm'])
    net['conv4_3_norm_mbox_conf_flat']  = Flatten(name='conv4_3_norm_mbox_conf_flat')(net['conv4_3_norm_mbox_conf'])

    # 对fc7层进行处理
    # 19,19,1024
    num_anchors = 6
    # 预测框的处理
    # num_anchors表示每个网格点先验框的数量，4是x,y,h,w的调整
    net['fc7_mbox_loc']         = Conv2D(num_anchors * 4, kernel_size=(3,3),padding='same',name='fc7_mbox_loc')(net['fc7'])
    net['fc7_mbox_loc_flat']    = Flatten(name='fc7_mbox_loc_flat')(net['fc7_mbox_loc'])
    # num_anchors表示每个网格点先验框的数量，num_classes是所分的类
    net['fc7_mbox_conf']        = Conv2D(num_anchors * num_classes, kernel_size=(3,3),padding='same',name='fc7_mbox_conf')(net['fc7'])
    net['fc7_mbox_conf_flat']   = Flatten(name='fc7_mbox_conf_flat')(net['fc7_mbox_conf'])

    # 对conv6_2进行处理
    # 10,10,512
    num_anchors = 6
    # 预测框的处理
    # num_anchors表示每个网格点先验框的数量，4是x,y,h,w的调整
    net['conv6_2_mbox_loc']         = Conv2D(num_anchors * 4, kernel_size=(3,3), padding='same',name='conv6_2_mbox_loc')(net['conv6_2'])
    net['conv6_2_mbox_loc_flat']    = Flatten(name='conv6_2_mbox_loc_flat')(net['conv6_2_mbox_loc'])
    # num_anchors表示每个网格点先验框的数量，num_classes是所分的类
    net['conv6_2_mbox_conf']        = Conv2D(num_anchors * num_classes, kernel_size=(3,3), padding='same',name='conv6_2_mbox_conf')(net['conv6_2'])
    net['conv6_2_mbox_conf_flat']   = Flatten(name='conv6_2_mbox_conf_flat')(net['conv6_2_mbox_conf'])

    # 对conv7_2进行处理
    # 5,5,256
    num_anchors = 6
    # 预测框的处理
    # num_anchors表示每个网格点先验框的数量，4是x,y,h,w的调整
    net['conv7_2_mbox_loc']         = Conv2D(num_anchors * 4, kernel_size=(3,3), padding='same',name='conv7_2_mbox_loc')(net['conv7_2'])
    net['conv7_2_mbox_loc_flat']    = Flatten(name='conv7_2_mbox_loc_flat')(net['conv7_2_mbox_loc'])
    # num_anchors表示每个网格点先验框的数量，num_classes是所分的类
    net['conv7_2_mbox_conf']        = Conv2D(num_anchors * num_classes, kernel_size=(3,3), padding='same',name='conv7_2_mbox_conf')(net['conv7_2'])
    net['conv7_2_mbox_conf_flat']   = Flatten(name='conv7_2_mbox_conf_flat')(net['conv7_2_mbox_conf'])

    # 对conv8_2进行处理
    # 3,3,256
    num_anchors = 4
    # 预测框的处理
    # num_anchors表示每个网格点先验框的数量，4是x,y,h,w的调整
    net['conv8_2_mbox_loc']         = Conv2D(num_anchors * 4, kernel_size=(3,3), padding='same',name='conv8_2_mbox_loc')(net['conv8_2'])
    net['conv8_2_mbox_loc_flat']    = Flatten(name='conv8_2_mbox_loc_flat')(net['conv8_2_mbox_loc'])
    # num_anchors表示每个网格点先验框的数量，num_classes是所分的类
    net['conv8_2_mbox_conf']        = Conv2D(num_anchors * num_classes, kernel_size=(3,3), padding='same',name='conv8_2_mbox_conf')(net['conv8_2'])
    net['conv8_2_mbox_conf_flat']   = Flatten(name='conv8_2_mbox_conf_flat')(net['conv8_2_mbox_conf'])

    # 对conv9_2进行处理
    # 1,1,256
    num_anchors = 4
    # 预测框的处理
    # num_anchors表示每个网格点先验框的数量，4是x,y,h,w的调整
    net['conv9_2_mbox_loc']         = Conv2D(num_anchors * 4, kernel_size=(3,3), padding='same',name='conv9_2_mbox_loc')(net['conv9_2'])
    net['conv9_2_mbox_loc_flat']    = Flatten(name='conv9_2_mbox_loc_flat')(net['conv9_2_mbox_loc'])
    # num_anchors表示每个网格点先验框的数量，num_classes是所分的类
    net['conv9_2_mbox_conf']        = Conv2D(num_anchors * num_classes, kernel_size=(3,3), padding='same',name='conv9_2_mbox_conf')(net['conv9_2'])
    net['conv9_2_mbox_conf_flat']   = Flatten(name='conv9_2_mbox_conf_flat')(net['conv9_2_mbox_conf'])

    # 将所有结果进行堆叠
    net['mbox_loc'] = Concatenate(axis=1, name='mbox_loc')([net['conv4_3_norm_mbox_loc_flat'],
                                                            net['fc7_mbox_loc_flat'],
                                                            net['conv6_2_mbox_loc_flat'],
                                                            net['conv7_2_mbox_loc_flat'],
                                                            net['conv8_2_mbox_loc_flat'],
                                                            net['conv9_2_mbox_loc_flat']])

    net['mbox_conf'] = Concatenate(axis=1, name='mbox_conf')([net['conv4_3_norm_mbox_conf_flat'],
                                                            net['fc7_mbox_conf_flat'],
                                                            net['conv6_2_mbox_conf_flat'],
                                                            net['conv7_2_mbox_conf_flat'],
                                                            net['conv8_2_mbox_conf_flat'],
                                                            net['conv9_2_mbox_conf_flat']])
    # 8732,4
    net['mbox_loc']     = Reshape((-1, 4), name='mbox_loc_final')(net['mbox_loc'])
    # 8732,21
    net['mbox_conf']    = Reshape((-1, num_classes), name='mbox_conf_logits')(net['mbox_conf'])
    net['mbox_conf']    = Activation('softmax', name='mbox_conf_final')(net['mbox_conf'])
    # 8732,25
    net['predictions']  = Concatenate(axis =-1, name='predictions')([net['mbox_loc'], net['mbox_conf']])

    model = Model(net['input'], net['predictions'])
    return model


import requests
import json

if __name__ == "__main__":
  # Downloads and patches the latest version of the image file

  # Download the source code at https://novelai.net/image without authentication to
  # find the link to the current image javascript file
  r = requests.get("https://novelai.net/image")
  image_page = r.content.decode("utf-8")
  image_script_url = "https://novelai.net/_next/static/chunks/pages/image-" + image_page.split('<script src="/_next/static/chunks/pages/image-')[-1].split('"')[0]
  print(f"Found current image script url: {image_script_url}")

  # Update the rules.json file to filter to the latest image file
  # An explicit version redirect is used so that if NAI updates
  # then no redirection will occur. NAI will still work, but will
  # not include a couple of these added features.
  with open("./../rules.json", "r") as f:
    rules = json.load(f)
  rules[0]["condition"]["urlFilter"] = image_script_url
  with open("./../rules.json", "w") as f:
    json.dump(rules, f, indent=2)
  print(f"Updated rules.json to filter to {image_script_url}")

  r = requests.get(image_script_url)
  image_script = r.content.decode("utf-8")
  with open("./../image_unpatched.js", "wb") as f:
    f.write(image_script.encode("utf-8"))
    print(f"Downloaded image file to image_unpatched.js")

  # Build the new image sizes to add
  new_image_sizes = """
      {name:"Portrait",width:${extra_small_short},height:${extra_small_long},category:"Extra Small"},
      {name:"Landscape",width:${extra_small_long},height:${extra_small_short},category:"Extra Small"},
      {name:"Square",width:${extra_small_square},height:${extra_small_square},category:"Extra Small"},
  """
  from string import Template
  new_image_sizes = Template(new_image_sizes).substitute(
    extra_small_square=str(64*6), # 384
    extra_small_long=str(64*8), # 512
    extra_small_short=str(int(64*5)), # 320
  )

  patches = [
    # Update max image count to 10
    # Note: Disabled due to NovelAI team not wanting to allow this feature.
    #("{var i=4;return e*t>", "{var i=10;return e*t>"),
    # Add "Extra Small" image size
    (
      '{name:"Square",width:512,height:512,category:"Small"},',
      '{name:"Square",width:512,height:512,category:"Small"},' + new_image_sizes
    )
  ]

  # Apply the patches
  for patch in patches:
    if patch[0] in image_script:
      print(f"Patching {patch[0]} -> {patch[1]}...")
      image_script = image_script.replace(patch[0], patch[1])
    else:
      print(f"ERROR: Patch {patch[0]} not found!")

  # Write the patched image file
  with open("./../image_patched.js", "wb") as f:
    f.write(image_script.encode("utf-8"))

  print(f"Done! Patched image file written to image_patched.js")

  # Beautify the patched image javascript file
  # NOTE: You'll need to install js-beautify via `npm install -g js-beautify` first
  print(f"Beautifying patched javascript...")
  import os
  os.system("js-beautify -r ./../image_patched.js") # Beautify the file in place
  print(f"Done! Beautified script overwritten to image_patched.js")