_SPLITS = {}
#################
# CTMCV1
#################
# sequences used for training
ctmcv1_train_seqs = ['3T3-run01']#['A-10-run01', '3T3-run01']#'3T3-run01','APM-run01','BPAE-run05','CV-1-run03','LLC-MK2-run07'] #, 'ETH-Sunnyday', 'ETH-Bahnhof', 'PETS09-S2L1', 'TUD-Stadtmitte']
all_seqs = ['3T3-run01','A-10-run01'] #,'APM-run01','BPAE-run05','CV-1-run03','LLC-MK2-run07','MDBK-run09','MDOK-run09','PL1Ut-run01','U2O-S-run03','3T3-run03','A-10-run03','APM-run03','BPAE-run07','LLC-MK2-run01','MDBK-run01','MDOK-run01','OK-run01','PL1Ut-run03','U2O-S-run05','3T3-run05','A-10-run05','APM-run05','CRE-BAG2-run01','LLC-MK2-run02a','MDBK-run03','MDOK-run03','OK-run03','PL1Ut-run05','3T3-run07','A-10-run07','BPAE-run01','CRE-BAG2-run03','LLC-MK2-run03','MDBK-run05','MDOK-run05','OK-run05','RK-13-run01','3T3-run09','A-549-run03','BPAE-run03','CV-1-run01','LLC-MK2-run05','MDBK-run07','MDOK-run07','OK-run07','RK-13-run03']

# Additional train sequences not used for tranining (since they are present in MOT17 etc.)
add_ctmcv1_train_seqs = []
# _SPLITS['mot15_train_gt'] = {'2DMOT2015/train': [f'{seq}-GT' for seq in mot15_train_seqs]}
_SPLITS['ctmcv1_train_gt'] = {'CTMCCVPR20/train': ctmcv1_train_seqs + add_ctmcv1_train_seqs}

# _SPLITS['ctmcv1_split_1_train_gt'] = {'CTMCCVPR20/train': ctmcv1_train_seqs}
# _SPLITS['split_1_ctmc_val'] = {'CTMCCVPR20/test': ['A-10-run02']}
# _SPLITS['ctmcv1_split_2_train_gt'] = {'CTMCCVPR20/train': ['A-10-run01']}
# _SPLITS['split_2_ctmc_val'] = {'CTMCCVPR20/train': ['3T3-run01']}

# This is used for debugging
#################################################################################
#_SPLITS['ctmcv1_split_1_train_gt'] = {'CTMCCVPR20/train': ['LLC-MK2-run07']}    #
#_SPLITS['split_1_ctmc_val'] = {'CTMCCVPR20/train': ['3T3-run01']}               #
#################################################################################
# _SPLITS['ctmcv1_split_1_train_gt'] = {'CTMCCVPR20/train': [f'{all_seqs[i]}' for i in range(len(all_seqs)) if i%2 == 0]}
# _SPLITS['split_1_ctmc_val'] = {'CTMCCVPR20/train': [f'{all_seqs[i]}' for i in range(len(all_seqs)) if i%2 == 1]}
# # _SPLITS['split_1_ctmc_val'] = {'CTMCCVPR20/train': [f'{all_seqs[i]}' for i in range(len(all_seqs)) if i%2 == 1]}
# _SPLITS['ctmcv1_split_2_train_gt'] = {'CTMCCVPR20/train': [f'{all_seqs[i]}' for i in range(len(all_seqs)) if i%2 == 1]}
# _SPLITS['split_2_ctmc_val'] = {'CTMCCVPR20/train': [f'{all_seqs[i]}' for i in range(len(all_seqs)) if i%2 == 0]}
# Test sequences
# test_seqs =  ['TUD-Crossing'] #, 'PETS09-S2L2', 'ETH-Jelmoli', 'ETH-Linthescher', 'ETH-Crossing', 'AVG-TownCentre',
#                 #   'ADL-Rundle-1', 'ADL-Rundle-3', 'KITTI-16', 'KITTI-19', 'Venice-1']
# _SPLITS['mot15_test'] = {'2DMOT2015/test': test_seqs}

# strat_over_framenum = [['MDBK-run09', 'MDOK-run09', 'LLC-MK2-run03', 'MDOK-run07', 'U2O-S-run05', 'U2O-S-run03', 'MDOK-run05', 'PL1Ut-run03', 'BPAE-run03', 'CRE-BAG2-run03', 'A-10-run05', 'APM-run03', 'OK-run01', 'MDBK-run01'], ['A-10-run07', '3T3-run03', '3T3-run07', 'CRE-BAG2-run01', 'LLC-MK2-run01', 'MDBK-run07', '3T3-run01', 'CV-1-run03', 'BPAE-run01', 'OK-run03', 'OK-run05'], ['BPAE-run07', '3T3-run09', '3T3-run05', 'A-10-run01', 'A-10-run03', 'LLC-MK2-run05', 'APM-run01', 'OK-run07', 'A-549-run03', 'RK-13-run01', 'MDBK-run05', 'RK-13-run03', 'MDBK-run03'], ['LLC-MK2-run07', 'MDOK-run01', 'APM-run05', 'CV-1-run01', 'BPAE-run05', 'PL1Ut-run01', 'PL1Ut-run05', 'LLC-MK2-run02a', 'MDOK-run03']]
# _SPLITS['ctmcv1_split_1_train_gt'] = {'CTMCCVPR20/train': strat_over_framenum[1]+strat_over_framenum[2]+strat_over_framenum[3]}
# _SPLITS['split_1_ctmc_val'] = {'CTMCCVPR20/train': strat_over_framenum[0]}
# _SPLITS['ctmcv1_split_2_train_gt'] = {'CTMCCVPR20/train': strat_over_framenum[0]+strat_over_framenum[2]+strat_over_framenum[3]}
# _SPLITS['split_2_ctmc_val'] = {'CTMCCVPR20/train': strat_over_framenum[1]}
# _SPLITS['ctmcv1_split_3_train_gt'] = {'CTMCCVPR20/train': strat_over_framenum[0]+strat_over_framenum[1]+strat_over_framenum[3]}
# _SPLITS['split_3_ctmc_val'] = {'CTMCCVPR20/train': strat_over_framenum[2]}
# _SPLITS['ctmcv1_split_4_train_gt'] = {'CTMCCVPR20/train': strat_over_framenum[0]+strat_over_framenum[1]+strat_over_framenum[2]}
# _SPLITS['split_4_ctmc_val'] = {'CTMCCVPR20/train': strat_over_framenum[3]}

divided_split_2_framenum_1 = [['LLC-MK2-run07',
  '3T3-run09',
  'APM-run05',
  'A-10-run01',
  'CV-1-run01',
  'BPAE-run05',
  'CV-1-run03',
  'A-10-run05',
  'MDBK-run05',
  'RK-13-run01',
  'LLC-MK2-run02a',
  'MDOK-run03'],
 ['MDBK-run09',
  'MDOK-run09',
  'LLC-MK2-run03',
  'A-10-run03',
  'LLC-MK2-run05',
  'BPAE-run03',
  'PL1Ut-run01',
  'APM-run03',
  'OK-run01']]

strat_over_framenum = [['A-10-run07', 'BPAE-run07', 'MDOK-run01', '3T3-run03', '3T3-run05', 'MDOK-run07', 'U2O-S-run05', 'U2O-S-run03', '3T3-run07', 'CRE-BAG2-run01', 'MDOK-run05', 'LLC-MK2-run01', 'PL1Ut-run03', 'MDBK-run07', 'APM-run01', '3T3-run01', 'CRE-BAG2-run03', 'OK-run07', 'A-549-run03', 'BPAE-run01', 'PL1Ut-run05', 'OK-run03', 'OK-run05', 'RK-13-run03', 'MDBK-run03', 'MDBK-run01'], ['MDBK-run09', 'LLC-MK2-run07', '3T3-run09', 'MDOK-run09', 'LLC-MK2-run03', 'APM-run05', 'A-10-run01', 'A-10-run03', 'LLC-MK2-run05', 'CV-1-run01', 'BPAE-run03', 'BPAE-run05', 'PL1Ut-run01', 'CV-1-run03', 'A-10-run05', 'RK-13-run01', 'MDBK-run05', 'APM-run03', 'OK-run01', 'LLC-MK2-run02a', 'MDOK-run03']]
_SPLITS['ctmcv1_split_1_train_gt'] = {'CTMCCVPR20/train': strat_over_framenum[1]}
_SPLITS['split_1_ctmc_val'] = {'CTMCCVPR20/train': strat_over_framenum[0]}
_SPLITS['ctmcv1_split_2_train_gt'] = {'CTMCCVPR20/train': strat_over_framenum[0]}
_SPLITS['split_2_ctmc_val'] = {'CTMCCVPR20/train': strat_over_framenum[1]}

#################
# MOT15
#################
dets = ('DPM', 'FRCNN', 'SDP')

# Train sequences:
mot15_train_seqs = ['KITTI-17' , 'ETH-Sunnyday', 'ETH-Bahnhof', 'PETS09-S2L1', 'TUD-Stadtmitte']

# Additional train sequences not used for tranining (since they are present in MOT17 etc.)
add_mot15_train_seqs = ['ETH-Pedcross2', 'TUD-Campus', 'KITTI-13', 'Venice-2', 'ADL-Rundle-8', 'ADL-Rundle-6']
_SPLITS['mot15_train_gt'] = {'2DMOT2015/train': [f'{seq}-GT' for seq in mot15_train_seqs]}
_SPLITS['mot15_train'] = {'2DMOT2015/train': mot15_train_seqs + add_mot15_train_seqs}

# Test sequences
test_seqs =  ['TUD-Crossing' , 'PETS09-S2L2', 'ETH-Jelmoli', 'ETH-Linthescher', 'ETH-Crossing', 'AVG-TownCentre',
                 'ADL-Rundle-1', 'ADL-Rundle-3', 'KITTI-16', 'KITTI-19', 'Venice-1']
_SPLITS['mot15_test'] = {'2DMOT2015/test': test_seqs}

#################
# MOT17
#################
dets = ('DPM', 'FRCNN', 'SDP')

# Train sequences:
train_seq_nums=  (2, 4, 5, 9, 10, 11, 13)
_SPLITS['mot17_train_gt'] = {'MOT17Labels/train': [f'MOT17-{seq_num:02}-GT' for seq_num in train_seq_nums]}
_SPLITS['mot17_train'] = {'MOT17Labels/train': [f'MOT17-{seq_num:02}-{det}' for seq_num in train_seq_nums for det in dets]}
_SPLITS['mot17_train_sdp'] = {'MOT17Labels/train': [f'MOT17-{seq_num:02}-SDP' for seq_num in train_seq_nums ]}

# Cross Validation splits
_SPLITS['mot17_split_1_train_gt'] = {'MOT17Labels/train': [f'MOT17-{seq_num:02}-GT' for seq_num in [2]]}
_SPLITS['split_1_val'] = {'MOT17Labels/train': [f'MOT17-{seq_num:02}-{det}' for seq_num in [4] for det in dets]}

# _SPLITS['mot17_split_1_train_gt'] = {'MOT17Labels/train': [f'MOT17-{seq_num:02}-GT' for seq_num in (2, 5, 9, 10, 13)]}
# _SPLITS['split_1_val'] = {'MOT17Labels/train': [f'MOT17-{seq_num:02}-{det}' for seq_num in (4,11) for det in dets]}

_SPLITS['mot17_split_2_train_gt'] = {'MOT17Labels/train': [f'MOT17-{seq_num:02}-GT' for seq_num in (2, 4, 11, 10, 13)]}
_SPLITS['split_2_val'] = {'MOT17Labels/train': [f'MOT17-{seq_num:02}-{det}' for seq_num in (5, 9) for det in dets]}

_SPLITS['mot17_split_3_train_gt'] = {'MOT17Labels/train': [f'MOT17-{seq_num:02}-GT' for seq_num in (4, 5, 9, 11)]}
_SPLITS['split_3_val'] = {'MOT17Labels/train': [f'MOT17-{seq_num:02}-{det}' for seq_num in (2, 10, 13) for det in dets]}

_SPLITS['debug'] = {'MOT17Labels/train': ['MOT17-02-FRCNN']}

# Test sequences
test_seq_nums=  (1, 3, 6, 7, 8, 12, 14)
_SPLITS['mot17_test'] = {'MOT17Labels/test': [f'MOT17-{seq_num:02}-{det}' for seq_num in test_seq_nums for det in dets]}

############
# MOT20
############

train_seq_nums=  (1, 2, 3, 5)
# Train / Val sequences
_SPLITS['mot20_train'] = {'MOT20/train': [f'MOT20-{seq_num:02}' for seq_num in train_seq_nums]}
_SPLITS['mot20_train_gt'] = {'MOT20/train': [f'MOT20-{seq_num:02}-GT' for seq_num in train_seq_nums]}
_SPLITS['mot20_train_wo_val'] = {'MOT20/train': [f'MOT20-{seq_num:02}-GT' for seq_num in (1, 2,  5)]}
_SPLITS['mot20_val'] = {'MOT20/train': [f'MOT20-{seq_num:02}' for seq_num in (3,)]}

_SPLITS['mot20_train_gt+'] = {'MOT17Labels/train': [f'MOT17-{seq_num:02}-GT' for seq_num in (2, 4, 9)],
                              'MOT20/train': [f'MOT20-{seq_num:02}-GT' for seq_num in train_seq_nums]}

# Cross-Val
for split_num, val_seq in enumerate(train_seq_nums, 1):
    _SPLITS[f'mot20_train_{split_num}'] = {'MOT17Labels/train': [f'MOT17-{seq_num:02}-GT' for seq_num in (2, 4, 9)],
                                'MOT20/train': [f'MOT20-{seq_num:02}-GT' for seq_num in train_seq_nums if seq_num != val_seq]}
    _SPLITS[f'mot20_val_{split_num}'] = {'MOT20/train': [f'MOT20-{val_seq:02}']}

# Test Sequences
_SPLITS['mot20_test'] = {'MOT20/test': [f'MOT20-{seq_num:02}' for seq_num in (4, 6, 7, 8)]}

# Combinations:
_SPLITS['all_train'] = {**_SPLITS['mot17_train_gt'], **_SPLITS['mot15_train_gt'], **_SPLITS['mot20_train_gt']}
_SPLITS['all_test'] = {**_SPLITS['mot17_test'], **_SPLITS['mot15_test'], **_SPLITS['mot20_test']}


import os

import torch
from torch import nn
import torchvision.models as torchvision_models

import tov_vicreg.models.networks.vit as vits
from tov_vicreg.models.networks.resnet import ResnetCNN
import tov_vicreg.utils.pytorch_utils as pytorch_utils

def load_model(arch, pretrained_weights=None, patch_size=8, num_classes=0, freeze=False, n_channels=3, input_shape=(1, 3, 84, 84)):
    if arch.startswith("vit_"):
        model = vits.__dict__[arch](patch_size=patch_size, num_classes=num_classes)
    elif arch == "sgi_resnet":
        model = ResnetCNN(input_channels=n_channels, depths=[64, 128, 256, 512], strides=[2,2,2,2])
    elif arch == "rainbow_cnn" or arch == "canonical":
        model = nn.Sequential(
            nn.Conv2d(n_channels, 32, 8, stride=4, padding=0),
            nn.ReLU(),
            nn.Conv2d(32, 64, 4, stride=2, padding=0),
            nn.ReLU(),
            nn.Conv2d(64, 64, 3, stride=1, padding=0),
            nn.ReLU(),
            nn.Flatten()
        )
    elif arch == "data-efficient" or arch == "der":
        model = nn.Sequential(
            nn.Conv2d(n_channels, 32, 5, stride=5, padding=0),
            nn.ReLU(),
            nn.Conv2d(32, 64, 5, stride=5, padding=0),
            nn.ReLU(),
            nn.Flatten()
        )
    elif arch == "linear":
        model = nn.Sequential(nn.Linear(4, 128), nn.ReLU(),)
        model.num_features = 128
    else:
        model = torchvision_models.__dict__[arch]()

    if "vit" not in arch or "linear" not in arch:
        with torch.no_grad():
            model.num_features = model(torch.ones(input_shape)).shape[1]

    for p in model.parameters():
        p.requires_grad = not freeze

    if pretrained_weights is None or pretrained_weights == "_":
        print(f"No pretrained weights provided, using random {arch}")
        return model

    if os.path.isfile(pretrained_weights):
        state_dict = torch.load(pretrained_weights, map_location="cpu")

        if "teacher" in state_dict.keys():  # DINO
            state_dict = state_dict["teacher"]
        elif "state_dict" in state_dict.keys():  # MOCOv3
            state_dict = state_dict["state_dict"]
            for k in list(state_dict.keys()):
                # retain only base_encoder up to before the embedding layer
                if k.startswith("base_encoder") and not k.startswith(
                    "base_encoder.head"
                ):
                    # remove prefix
                    state_dict[k[len("base_encoder.") :]] = state_dict[k]
                # delete renamed or unused k
                del state_dict[k]
        elif list(state_dict.keys())[0].startswith("vit."):
            for k in list(state_dict.keys()):
                if k.startswith("vit") and not k.startswith("vit.head"):
                    # remove prefix
                    state_dict[k[len("vit.") :]] = state_dict[k]
                # delete renamed or unused k
                del state_dict[k]
        elif list(state_dict.keys())[0].startswith("encoder."):
            for k in list(state_dict.keys()):
                if k.startswith("encoder") and not k.startswith("encoder.head"):
                    # remove prefix
                    state_dict[k[len("encoder.") :]] = state_dict[k]
                # delete renamed or unused k
                del state_dict[k]

        # remove `module.` prefix
        state_dict = {k.replace("module.", ""): v for k, v in state_dict.items()}
        # remove `backbone.` prefix induced by multicrop wrapper
        state_dict = {k.replace("backbone.", ""): v for k, v in state_dict.items()}
        # remove head
        for k in list(state_dict.keys()):
            if k.startswith("head."):
                del state_dict[k]
        msg = model.load_state_dict(state_dict, strict=False)
        print(
            "Pretrained weights found at {} and loaded with msg: {}".format(
                pretrained_weights, msg
            )
        )
    else:
        raise FileExistsError("Can't find file with pretrained weights")
    return model


import linearbag

myBag = linearbag.Bag()
m_lst = [10, 20, 30, 40]
myBag.add(10)
myBag.add(20)
myBag.add(30)
myBag.add(40)

print(myBag.__iter__(m_lst).__dict__)
print(myBag._items)

value = int( input("Guess a value contained in the bag: ") )
if value in myBag:
    print( "The bag contains the value", value )
else :
    print( "The bag does not contain the value", value )


import timeit

numbers = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
even_numbers = [x for x in numbers if x % 2 == 0]
print(even_numbers)

numbers = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
even_numbers = (x for x in numbers if x % 2 == 0)
# print(next(even_numbers))
# print(next(even_numbers))

# using list comprehension
def using_list_comprehension():
    even_numbers = [x for x in numbers if x % 2 == 0]

# using a generator expression
def using_generator_expression():
    even_numbers = (x for x in numbers if x % 2 == 0)

print("Using list comprehension:", timeit.timeit(using_list_comprehension, number=1000))
print(
    "Using generator expression:",
    timeit.timeit(using_generator_expression, number=1000)
)


import streamlit as st
import pandas as pd
import numpy as np

import gensim.downloader as api
from gensim.models import KeyedVectors, Word2Vec

import plotly.graph_objs as go
from sklearn.decomposition import PCA

st.title('Visualizing Word Embeddings')

model_wikipedia50 = api.load("glove-wiki-gigaword-50")
data = model_wikipedia50

st.header('What Are Word Embeddings?')
st.write('Word embeddings are vector representations of words. Words with similar meanings are closer together in vector space.')
st.write('Search for any word below and the graph will show the word embeddings of the most similar words!')

title = st.text_input('Word (Examples: "archaeologists", "baseball", "Google", "2006", "meningococcus", "Ushuaia")', 'dog').lower()

if st.button('Random Word'):
    title = model_wikipedia50.index_to_key[np.random.randint(0,399999)]

st.write('Selected Word: ', title)

num = st.slider('Select number of similar words:', 5, 50, 10)

def append_list(sim_words, words):

    list_of_words = []

    for i in range(len(sim_words)):

        sim_words_list = list(sim_words[i])
        sim_words_list.append(words)
        sim_words_tuple = tuple(sim_words_list)
        list_of_words.append(sim_words_tuple)

    return list_of_words

try:
    word_list = [model_wikipedia50.most_similar(f'{title}', topn=num)[i][0] for i in range(num)]
    words = str(word_list).replace("'","")
    words = words.strip("[]")
    data=[model_wikipedia50.most_similar(f'{title}', topn=num)[i] for i in range(num)]
    df = pd.DataFrame(data=data, columns= ['Word','Cosine Similarity'])
    df.index +=1

    input_word = title.lower() + ',' + words
    user_input = [x.strip() for x in input_word.split(',')]
    result_word = []

    for words in user_input:

            sim_words = model_wikipedia50.most_similar(words, topn = 0)
            sim_words = append_list(sim_words, words)

            result_word.extend(sim_words)

    similar_word = [word[0] for word in result_word]
    similarity = [word[1] for word in result_word]
    similar_word.extend(user_input)
    labels = [word[2] for word in result_word]
    label_dict = dict([(y,x+1) for x,y in enumerate(set(labels))])
    color_map = [label_dict[x] for x in labels]

except KeyError:
    st.write('Sorry! \"',title,"\" not in vocabulary." )

def display_pca_scatterplot_3D(model=model_wikipedia50, user_input=None, words=None, label=None, color_map=None, topn=5, sample=10):

    if words == None:
        if sample > 0:
            words = np.random.choice(list(model.vocab.keys()), sample)
        else:
            words = [ word for word in model.vocab ]

    word_vectors = np.array([model[w] for w in words])

    three_dim = PCA(random_state=0).fit_transform(word_vectors)[:,:3]

    data = []
    count = 0

    for i in range (len(user_input)):
        trace = go.Scatter3d(
        x = three_dim[count:count+topn,0],
        y = three_dim[count:count+topn,1],
        z = three_dim[count:count+topn,2],
        text = words[count:count+topn],
        name = user_input[i],
        textposition = "top center",
        textfont_size = 20,
        mode = 'markers+text',
        marker = {
            'size': 10,
            'opacity': 0.8,
            'color': 2
        }

                )

        data.append(trace)
        count = count+topn

    trace_input = go.Scatter3d(
                    x = three_dim[count:,0],
                    y = three_dim[count:,1],
                    z = three_dim[count:,2],
                    text = words[count:],
                    name = 'input words',
                    textposition = "top center",
                    textfont_size = 20,
                    mode = 'markers+text',
                    marker = {
                        'size': 10,
                        'opacity': 1,
                        'color': 'black'
                    }
                    )

    data.append(trace_input)

    layout = go.Layout(
        margin = {'l': 0, 'r': 0, 'b': 0, 't': 0},
        showlegend=False,
        legend=dict(
        x=1,
        y=0.5,
        font=dict(
            family="Courier New",
            size=25,
            color="black"
        )),
        font = dict(
            family = " Courier New ",
            size = 15),
        autosize = False,
        width = 700,
        height = 700
        )

    plot_figure = go.Figure(data = data, layout = layout)
    st.plotly_chart(plot_figure)

try: display_pca_scatterplot_3D(model_wikipedia50, user_input, similar_word, labels, color_map)
except NameError:
    pass

try:
    st.write(num, 'Most Similar Words:')
    st.table(df)
except NameError:
    pass


import os
import sys
import torch
import tqdm
import json
import numpy as np
from constree    import *
from billion_words import *
from collections import Counter
from random      import shuffle,random
import torch.nn as nn
import torch.optim as optim
from torch.nn.utils import clip_grad_norm_
from torch.nn.utils.rnn import pack_padded_sequence,pad_packed_sequence
from torch.optim.lr_scheduler import LambdaLR,ReduceLROnPlateau

#MODELS
#Derivation = w1 s1 w2 s2   ...   wN sN
#Model 1
#  P(wi | w1 s1 ... wi-1  si-1)
#Estimation
#    ~ P(w_i | w1 ... wi-1 ) [ task 1 : language model where structure is only used at estimation]
#      P(s_i | w1 ... w_i  ) [ task 2 : predict annotations outside language model. ]

#with the additional independance assumptions we have :
#P(w_i | w1 ... wi-1 ) decomposes as:
#   P(wordlabel_i | w1 ... wi-1 ) P(wordaction_i | w1 ... wi-1)
#P(s_i | w1 ... wi ) decomposes as:
#   P(NTlabel_i | w1 ... wi ) P(NTaction_i | w1 ... wi)

#The key property of the model is that it can learn from unsupervised data the language model subtask with
#   >>> P(wordlabel_i | w1 ... wi-1 ) <<<

# Prediction
# these substasks may use constraints that create a conditioning on previous hidden struct too -> not all struct actions are always possible
#use beam search
# Exact decomposition :
# prod_i^N  P(wi | w1,s1 ...  wi-1,si-1) P(si | (w1,s1) ... (wi-1,si-1), wi)
# Our approximation (in exchange for cheap unsup learning)  :
# prod_i^N  P(wi | w1 ...  wi-1)         P(si | w1 ... wi-1, wi)
# We add an additional conditional independance assumption to isolate the language model :
# xxx

#Model 2
#  P(ai | a1 ...  ai-1) -> requires true beam search because si gets back on input like in seq2seq
#                          (for training it is not sure that it changes a lot, requires to input the struct actions but that's it)
#  unsupervised = xxx boils down to basic domain adaptation

class LCtree (ConsTree):
    """
    Specialization of ConsTree for coding complete/incomplete constituents as used in this transition system.
    """
    def __init__(self,label,children=None):
        super(LCtree, self).__init__(label,children)

    def rightmost_attach_site(self):
        """
        Returns the rightmost node in a constituent tree
        to which we can attach a right child.
        Args:
             ctree (ConsTree): tree where to search
        Returns:
              the immediate root of the attachement site if such node exists else None
        """
        if self.arity() == 2:
            return self.get_child(1).rightmost_attach_site()
        else: #arity < 2
            return None if self.is_leaf() else self

    def is_complete(self):
        return self.rightmost_attach_site() is None

#DATA SET
#This recodes with more flexibility some key functionalities for managing data inspired by
#torch.text by making them suitable for use for parsing purposes.

class Vocabulary:

    def __init__(self,counter,min_freq=1,pad='<pad>',unk=None,sos=None):
        """
        New vocabulary.
        Args:
            counter (collections.Counter): tokens frequencies
            min_freq                (int): words with counts strictly less than min_freq are removed from  the vocabulary
            unk                  (string): the reserved 'unk' token
            pad                  (string): the reserved 'pad' token
            sos                  (string): the reserved start of sentence token
        """
        self.unk = unk
        self.pad = pad
        self.sos = sos
        self.counts = Counter(dict([ (item,count) for (item,count) in counter.items() if count >= min_freq] ))

        #codes the lexicon to make it compatible with the adaptive logsoftmax
        self.itos   = [ ]
        if unk and not unk in self.counts:
            self.itos.append( unk )
        if pad and not pad in self.counts:
            self.itos.append( pad )
        if sos and sos not in self.counts:
            self.itos.append( sos )
        self.itos.extend( [ elt for (elt, count) in self.counts.most_common() ]) # <- here
        self.stoi = dict( zip(self.itos,range(len(self.itos))))

    def save(self,filename):
        """
        Saves vocabulary to file
        """
        ostream = open(filename+'.specials','w')
        print(self.unk,file=ostream)
        print(self.pad,file=ostream)
        print(self.sos,file=ostream)
        ostream.close()

        ostream = open(filename+'.counts','w')
        for token,counts in self.counts.items():
            print("%s\t%d"%(token,counts),file=ostream)
        ostream.close()

    @staticmethod
    def load(filename):
        """
        Loads a vocabulary from file and returns it
        """
        istream     = open(filename+'.counts')
        counts = Counter()
        for line in istream:
            token, count = line.split()
            counts[token] = int(count)
        istream.close()

        istream   = open(filename+'.specials')
        unk  = istream.readline()[:-1]
        pad  = istream.readline()[:-1]
        sos  = istream.readline()[:-1]
        istream.close()
        if unk == 'None':
            unk = None
        if pad == 'None':
            pad = None
        if sos == 'None':
            sos = None
        return Vocabulary(counts,pad=pad,sos=sos,unk=unk,min_freq=0)

    def size(self):
        """
        Returns the number of symbols in this vocabulary
        """
        return len(self.itos)

    def token_index(self,token):
        """
        Gets the index of a token
        Args:
           token  (str): the token for which we seek the index
        Returns:
           the index of the token, the index of the unk if not found and an error if unk is not in this dict
        """
        return self.stoi.get(token,0) if self.unk else self.stoi[token]

    def token_form(self,tok_idx):
        """
        Retrieves a token string from its index.
        Note that if the token was unknown to the vocabulary, the orginal form is not retrieved (rather the unk code)
        Args:
           tok_idx (int): the token integer index
        Returns:
           a string.
        """
        return self.itos[tok_idx]

class ParsingDataSet(object):
    """
    That's a data set for parsing. Each example is a couple made of a list of tokens and an optional derivation.
    That's currently tied to the parser class (try to remove this dependency later)
    """
    def __init__(self,dataset,ext_vocab=None,root_dataset=None, unk='<unk>',pad='<pad>',sos='<sos>',min_lex_counts=0):
        """
        Args:
             dataset               (list): a list of trees (or a list of strings for test and unsupervised setups)
             ext_vocab       (Vocabulary): if specified use existing external vocabulary for lexical words only, otherwise builds it from the data
             root_dataset(ParsingDataSet): if specified uses the encodings from this external dataset rather than inferring encodings from this one.
             unk                    (str): a string for the unk token for internal vocab
             pad                    (str): a string for the pad token for internal vocab
             sos                    (str): a string for the start of sentence (sos) token
             min_lex_counts         (int): count threshold under which tokens are excluded from the lexical vocabulary built from the data
        """
        is_treebank = isinstance(dataset[0],ConsTree)

        #1. Data structuration
        if is_treebank : #has annotated trees

            self.tree_set      = dataset
            derivations        = [ LCmodel.oracle_derivation(tree) for tree in dataset ]

            self.tokens        = [ tree.tokens() for tree in dataset ]
            self.lex_actions   = [ self.extract_lex_actions(deriv)  for deriv in derivations ]
            self.struct_labels = [ self.extract_struct_labels(deriv)  for deriv in derivations ]
            self.struct_actions= [ self.extract_struct_actions(deriv)  for deriv in derivations ]

        else:                             #raw text
            self.tokens = [ sent.split( ) for sent in dataset ]

        #2. Vocabularies
        if root_dataset :
            self.lex_vocab = ext_vocab if ext_vocab else root_dataset.lex_vocab
            if is_treebank :
                self.lex_action_vocab    = root_dataset.lex_action_vocab
                self.struct_vocab        = root_dataset.struct_vocab
                self.struct_action_vocab = root_dataset.struct_action_vocab
            self.unk = root_dataset.unk
            self.pad = root_dataset.pad
            self.sos = root_dataset.sos
        else:
            self.lex_vocab = ext_vocab if ext_vocab else ParsingDataSet.build_vocab(self.tokens,unk_lex=unk,pad=pad,sos=sos,min_counts=min_lex_counts)
            if is_treebank :
                self.lex_action_vocab    = ParsingDataSet.build_vocab(self.lex_actions,pad=pad)
                self.struct_vocab        = ParsingDataSet.build_vocab(self.struct_labels,pad=pad,sos=sos)
                self.struct_action_vocab = ParsingDataSet.build_vocab(self.struct_actions,pad=pad,sos=sos)
            self.unk = unk
            self.pad = pad
            self.sos = sos

    @staticmethod
    def load_dataset(path):
        if os.path.isfile(os.path.join(path,'train.mrg')):
            p = ParsingDataSet(list(input_treebank(os.path.join(path,'train.mrg'))))
        else:
            istream = open(os.path.join(path,'train.raw'))
            p = ParsingDataSet([line for line in istream])
            istream.close()
        p.lex_vocab = Vocabulary.load(os.path.join(path,'lex_vocab'))
        p.struct_vocab = Vocabulary.load(os.path.join(path,'struct_vocab'))
        p.struct_action = Vocabulary.load(os.path.join(path,'struct_action_vocab'))
        p.lex_action_vocab = Vocabulary.load(os.path.join(path,'lex_action_vocab'))
        return p

    def save_dataset(self,path):
        """
        Saves the dataset to a set of files prefixed by path
        """
        try:
            os.mkdir(path)
        except :
            pass
        if hasattr(self,'tree_set'):
            ostream = open(os.path.join(path,'train.mrg'),'w')
            for tree in self.tree_set:
                print(str(tree),file=ostream)
            ostream.close()
        else:
            ostream = open(os.path.join(path,'train.raw'),'w')
            for toklist in self.tokens:
                print(' '.join(toklist),file=ostream)
            ostream.close()

        self.lex_vocab.save(os.path.join(path,'lex_vocab'))
        self.struct_vocab.save(os.path.join(path,'struct_vocab'))
        self.struct_action_vocab.save(os.path.join(path,'struct_action_vocab'))
        self.lex_action_vocab.save(os.path.join(path,'lex_action_vocab'))

    def decode_derivation(self,derivation): #pred_lexaction,pred_ytokens,pred_structaction,pred_structlabels
        """
        This translates an integer coded derivation back to a string.
        Args :
           xtokens    (list): list of integers (token codes)
           derivation (list): list of integer tuples
        """
        def translate(action,label,idx):
            if idx % 2 == 0 :
                return (self.lex_action_vocab.itos[action],self.lex_vocab.itos[label])
            else:
                return (self.struct_action_vocab.itos[action],self.struct_vocab.itos[label])

        return [translate(action,label,idx) for idx, (action,label) in enumerate(derivation) ]

    def is_training_set(self):
        """
        Returns true if this dataset has annotated trees
        """
        return hasattr(self,'lex_actions')

    @staticmethod
    def build_vocab(datalist,unk_lex=None,pad=None,sos=None,min_counts=0):
        """
        Builds a vocabulary from a matrix of strings
        Args:
           datalist (list of list) : a list of ConsTree objects
           unk_lex           (str) : a string to use for replacement of unknown words
           pad               (str) : a string to use padding tokens
           sos               (str) : a string to use start of sentence tokens
        Returns
            a Vocabulary object.
        """
        token_counter    = Counter()
        for item in datalist:
            token_counter.update(item)
        return Vocabulary(token_counter,unk=unk_lex,pad=pad,sos=sos,min_freq=min_counts)

    def extract_tokens(self,derivation):
        """
        Extracts the list of words from a derivation and returns it
        Args:
            derivation (list): a list of couples (action,label)
        Return :
            list. A list of word tokens
        """
        return [ label for (idx,(action,label)) in enumerate(derivation) if idx % 2 == 0]

    def extract_lex_actions(self,derivation):
        """
        Extracts the list of lexical actions from a derivation and returns it
        Args:
            derivation (list): a list of couples (action,label)
        Return :
            list. A list of word tokens
        """
        return [ action for (idx,(action,label)) in enumerate(derivation) if idx % 2 == 0]

    def extract_struct_labels(self,derivation):
        """
        Extracts the list of non terminal symbols from a derivation and returns it
        Args:
            derivation (list): a list of couples (action,label)
        Return :
            list. A list of non terminal symbols
        """
        return [ label for (idx,(action,label)) in enumerate(derivation) if idx % 2 == 1]

    def extract_struct_actions(self,derivation):
        """
        Extracts the list of structural actions from a derivation and returns it
        Args:
            derivation (list): a list of couples (action,label)
        Return :
            list. A list of structural actions
        """
        return [ action for (idx,(action,label)) in enumerate(derivation) if idx % 2 == 1]

    def __len__(self):
        """
        Returns the number of samples in the data set
        """
        return len(self.tokens)

    def example_length(self,idx):
        """
        Returns the length of an example (token length for batch sorting purposes)
        """
        return len(self.tokens[idx])

    def sample_tokens(self,toklist,batch_len,alpha=0.0):
        """
        This returns a list of integers encoding the tokens
        Args:
            toklist         (list): a list of strings
            alpha          (float): alpha >= 0 a constant used for sampling unk words. if alpha == 0, no sampling occurs
        """
        def sample_unk(token,alpha):
            prob = random()
            C = self.lex_vocab.counts[token]
            if C == 0 or prob < alpha / C:
                return self.unk
            return token

        def subst_unk(token):
            C = self.lex_vocab.counts[token]
            if C == 0:
                return self.unk
            return token

        if alpha > 0:
            return [sample_unk(tok,alpha) for tok in toklist ]
        else:
            return [subst_unk(tok) for tok in toklist ]

    def numericalize_example(self,datum,batch_len,vocabulary):
        """
        This returns a list of integers encoding the given symbols.
        Args:
            datum           (list): list of strings, the tokens to map to integers
            batch_len        (int): the length of the batch (>= actual number of tokens). Missing actions are padded left
            vocabulary(Vocabulary): a vocabulary object that performs the tokenwise mapping
        Returns:
            A list of integers
        """
        N                 = len(datum)
        padded_datum      = datum + [self.pad]*(batch_len-N)
        return [vocabulary.token_index(token) for token in padded_datum]

class ParseBatch:

    def __init__(self,xtokens=None,ytokens=None,struct_actions=None,struct_labels=None,lex_actions=None,lex_labels=None,token_length=None,orig_idxes=None,derivation_length=None):
        self.xtokens           = xtokens
        self.ytokens           = ytokens
        self.lex_actions       = lex_actions
        self.lex_labels        = lex_labels
        self.struct_actions    = struct_actions
        self.struct_labels     = struct_labels
        self.tokens_length     = token_length
        self.orig_idxes        = orig_idxes

class BucketLoader:
    """
    This BucketLoader is an iterator that reformulates torch.text.data.BucketIterator
    TODO: add multiprocessing capacities (load batch while computing another)
    """
    def __init__(self,dataset,batch_size,device=-1,alpha=0.0):
        """
        This function is responsible for delivering batches of examples from the dataset for a given epoch.
        It attempts to bucket examples of the same size in the same batches.

        Generated batches contain approximatively the same number of words.
        However the number of sentences in a batch may differ.
        If a single sentence contains more words than the batch size, it is simply dropped

        Args:
          dataset   (ParsingDataset): the treebank from where to get samples
          batch_size           (int): the size of the generated batches (in words !)
          device               (str): a string that says on which device the batched data should live (defaults to cpu).
          alpha              (float): param for sampling unk words
        """
        self.dataset      = dataset
        self.batch_size   = batch_size
        self.device       = device
        self.alpha        = alpha
        self.data_idxes   = list(range(len(self.dataset)))
        self.make_batches()

    def nbatches(self):
        return len(self.start_end_positions)

    def encode_batch(self,batch_idxes):
        """
        Encode and adds padding symbols to shorter examples in this batch
        Args:
            batch_idxes (list): a list of integers, indexes of the examples included in this batch
        Returns:
            a ParseBatch. The batch object fields are filled partly if not supervised training set
            The batches examples are padded right.
        """
        batchN             = len(batch_idxes)
        token_lengths      = [ len(self.dataset.tokens[idx]) for idx in batch_idxes ]
        max_token_length   = max(token_lengths)

        raw_tokens    = [ self.dataset.sample_tokens(self.dataset.tokens[batch_idxes[step]],max_token_length,alpha=self.alpha) for step in range(batchN) ]
        ytoken_matrix = [ self.dataset.numericalize_example(elt,max_token_length,self.dataset.lex_vocab) for elt in raw_tokens ]
        xtoken_matrix = [ self.dataset.numericalize_example([self.dataset.sos]+elt[:-1],max_token_length,self.dataset.lex_vocab) for elt in raw_tokens ]

        xtoken_tensor  = torch.tensor(xtoken_matrix,dtype=torch.long,device=self.device)
        ytoken_tensor  = torch.tensor(ytoken_matrix,dtype=torch.long,device=self.device)

        #print(raw_tokens)

        if self.dataset.is_training_set():

            lex_action_matrix     = [self.dataset.numericalize_example( self.dataset.lex_actions[batch_idxes[step]], max_token_length,self.dataset.lex_action_vocab)  for step in range(batchN) ]
            struct_action_matrix  = [self.dataset.numericalize_example([self.dataset.sos]+self.dataset.struct_actions[batch_idxes[step]], max_token_length,self.dataset.struct_action_vocab) for step in range(batchN) ]
            struct_label_matrix   = [self.dataset.numericalize_example([self.dataset.sos]+self.dataset.struct_labels[batch_idxes[step]], max_token_length,self.dataset.struct_vocab) for step in range(batchN) ]

            lex_action_tensor     = torch.tensor(lex_action_matrix,dtype=torch.long,device=self.device)
            struct_action_tensor  = torch.tensor(struct_action_matrix,dtype=torch.long,device=self.device)
            struct_label_tensor   = torch.tensor(struct_label_matrix,dtype=torch.long,device=self.device)

            return ParseBatch(xtokens=xtoken_tensor,ytokens=ytoken_tensor, lex_actions=lex_action_tensor,\
                                  struct_actions=struct_action_tensor,struct_labels=struct_label_tensor,token_length=token_lengths,orig_idxes=batch_idxes)

        return ParseBatch(xtokens=xtoken_tensor,ytokens=ytoken_tensor,token_length=token_lengths,orig_idxes=batch_idxes)

    def __iter__(self):
        return self

    def batch_range(self,r0):
        """
        Computes the number of sentences to include in this batch and returns a couple
        (p0,pe) of begin and end index position for this batch
        Args:
           p0        (int): begin position of this batch
        Returns
           a tuple (p0,pe). the begin and end position of this batch
           a tuple (p0,p0) in case no valid batch can be generated starting from p0
        """
        L0 = self.dataset.example_length(self.data_idxes[r0])
        N  = int(self.batch_size/L0)
        #print(N)
        return (r0,r0+N)

    def make_batches(self):
        """
        Actually makes the batches
        """
        shuffle(self.data_idxes)
        lengths              = [ self.dataset.example_length(idx) for idx in self.data_idxes ]
        self.data_idxes      = [ idx for (idx,length) in sorted(zip(self.data_idxes,lengths),key=lambda x:x[1],reverse=True) ]

        self.start_end_positions = [ ]
        cpos = 0
        while cpos < len(self.data_idxes):
            p0,pE = self.batch_range(cpos)
            if p0 == pE: #invalid batch of size 0
                cpos +=1
                print('invalid batch detected (sentence too long) skipped.',file=sys.stderr,flush=True)
            else:
                self.start_end_positions.append((p0,pE))
                cpos = pE
        shuffle(self.start_end_positions)

    def __next__(self):
        """
        This yields a batch of data.
        """
        if self.start_end_positions:
            start_pos,end_pos    = self.start_end_positions.pop()
            batch_idxes          = self.data_idxes[ start_pos:end_pos ]
            return self.encode_batch(batch_idxes)
        else:
            raise StopIteration

class LCmodel(nn.Module):

    ACTION_SHIFT_INIT    = 'I'
    ACTION_SHIFT_ATTACH  = 'S'
    ACTION_ATTACH        = 'A'
    ACTION_PREDICT       = 'P'

    def __init__(self,ref_set,rnn_memory=100,embedding_size=100,device=-1):
        """
        This allocates the model parameters on the machine.
        Args:
           ref_set (ParsingDataSet): the reference training from which vocabularies are built
           rnn_memory  (int)  : the size of the rnn hidden memory
           embedding_size(int): the embedding size
           device        (int): the device where to store the params (-1 :cpu ; 0,1,2... : GPU identifier)
        """
        super(LCmodel, self).__init__()
        self.ref_set        = ref_set
        self.rnn_memory     = rnn_memory
        self.embedding_size = embedding_size
        self.allocate_structure(device)
        if device >= 0:
            self.cuda(device=device)

    @staticmethod
    def load(path,device=-1):
        refset  = ParsingDataSet.load_dataset(os.path.join(path,'data'))
        istream = open(os.path.join(path,'hyperparams.json'))
        hparams = json.loads(istream.read())
        istream.close()
        m = LCmodel(refset,rnn_memory=hparams['rnn_hidden'],embedding_size=hparams['embedding_size'],device=device)
        m.load_state_dict(torch.load(os.path.join(path,'weights.wt')))
        return m

    def save(self,path):
        """
        Saves the model to prefix path
        """
        try:
            os.mkdir(path)
        except :
            pass
        self.ref_set.save_dataset(os.path.join(path,'data'))
        torch.save(self.state_dict(), os.path.join(path,'weights.wt'))
        ostream = open(os.path.join(path,'hyperparams.json'),'w')
        print(json.dumps({'rnn_hidden':self.rnn_memory,'embedding_size':self.embedding_size}),file=ostream)
        ostream.close()

    def allocate_structure(self,device=-1):
        """
        This allocates the model parameters on the machine.
        Args:
           action_size   (int): the number of action types
           lex_size      (int): the size of the lexicon vocabulary
           struct_size   (int): the size of the non terminal vocabulary
           rnn_memory    (int): the size of the rnn hidden memory
           embedding_size(int): the embedding size
           device        (int): the device where to store the params (-1 :cpu ; 0,1,2... : GPU identifier)
        """
        self.E               = nn.Embedding( self.ref_set.lex_vocab.size(),self.embedding_size)
        self.lstm            = nn.LSTM(self.embedding_size, self.rnn_memory,num_layers=1,bidirectional=False)

        self.W_struct_label  = nn.Linear(self.rnn_memory, self.ref_set.struct_vocab.size())
        self.W_lex_action    = nn.Linear(self.rnn_memory, self.ref_set.lex_action_vocab.size())
        self.W_struct_action = nn.Linear(self.rnn_memory, self.ref_set.struct_action_vocab.size())
        self.logsoftmax      = nn.LogSoftmax(dim=1)

        vs                   = self.ref_set.lex_vocab.size()
        self.adalogsoftmax   = nn.AdaptiveLogSoftmaxWithLoss(self.rnn_memory,vs,[ int(vs/15), 2*int(vs/15), 4*int(vs/15), 8*int(vs/15)], div_value=4.0, head_bias=False)

    def forward_lexical_actions(self,base_output,ref_actions=None,loss=None):
        """
        Performs the forward pass for the lexical action subtask
        Args:
            base_output  (tensor): the tensor outputted by the base LSTM encoder.
            ref_action   (tensor): the tensor with correct predictions
            loss   (pytorch loss): the function used to compute the loss
        Returns:
            a tensor. A list of softmaxed predictions for each example provided as argument.
            if ref_action and loss are provided, returns a tuple (softmaxed output ,loss)
        """
        output = self.logsoftmax(self.W_lex_action(base_output))

        if ref_action and loss:
            ref_lexactions     =  ref_actions.view(-1)
            L      = loss(output,ref_lexactions)
            return (output,L)
        return output

    def forward_lexical_tokens(self,base_output,ref_output):
        """
        Performs the forward pass for the lexical token subtask
        Args:
            base_output  (tensor): the tensor outputted by the base LSTM encoder.
            ref_output   (tensor): the tensor encoding the next word
        Returns:
             a tuple (softmaxed output ,loss). A list of softmaxed word predictions for each example provided as argument and the logsoftmax loss for ref_output
        """
        ref_output      =  ref_output.view(-1)

        #implements some sort of masking of the padded elts on the fly... probably easier ways to achieve that but the loss has no built-in.
        pad_idx         = self.ref_set.lex_vocab.stoi[self.ref_set.pad]
        nonpadded_idxes =  (ref_output != pad_idx).nonzero().t().squeeze()
        ref_output      =  ref_output[nonpadded_idxes]
        base_output     =  base_output[nonpadded_idxes]

        return self.adalogsoftmax(base_output,ref_output)

    def forward_structural_actions(self,base_output,ref_action=None,loss=None):
        """
        Performs the forward pass for the structural actions subtask
        Args:
            base_output  (tensor): the tensor outputted by the base LSTM encoder.
            ref_action   (tensor): the tensor with correct predictions
            loss   (pytorch loss): the function used to compute the loss
        Returns:
            a tensor. A list of softmaxed actions predictions for each example provided as argument.
            if ref_action and loss are provided, returns a tuple (softmaxed output ,loss)
        """
        output = self.logsoftmax(self.W_struct_action(base_output))
        if ref_action and loss:
            ref_action  =  ref_action.view(-1)
            L           = loss(output,ref_action)
            return (output,L)
        return output

    def forward_structural_labels(self,base_output,ref_action=None,loss=None):
        """
        Performs the forward pass for the structural labels subtask
        Args:
            base_output  (tensor): the tensor outputted by the base LSTM encoder.
            ref_action   (tensor): the tensor with correct predictions
            loss   (pytorch loss): the function used to compute the loss
        Returns:
            a tensor. A list of softmaxed non terminal predictions for each example provided as argument.
            if ref_action and loss are provided, returns a tuple (softmaxed output ,loss)
        """
        output = self.logsoftmax(self.W_struct_label(base_output))
        if ref_action and loss:
            ref_action  =  ref_action.view(-1)
            L           = loss(output,ref_action)
            return (output,L)
        return output

    def forward_base(self,xinput,true_batch_lengths):
        """
        Args :
           xinput           (tensor): an integer coded input
           true_batch_lengths (list): list of integers (true lengths of inputs)
        Returns:
           The RNN output as a pytorch flattened Sequence of wordwise encodings
        """
        #@see (packing) https://gist.github.com/HarshTrivedi/f4e7293e941b17d19058f6fb90ab0fec
        xembedded         = self.E(xinput)                                                        #xembedded [dim] = batch_size x sent_len x embedding_size
        xembedded         = pack_padded_sequence(xembedded, true_batch_lengths, batch_first=True)
        lstm_out, _       = self.lstm(xembedded)
        lstm_out, _       = pad_packed_sequence(lstm_out,batch_first=True)                        #lstm_out  [dim] = batch_size x sent_len x hidden_size

        # We flatten the batch before applying the linear outputs (required by the loss functions at the end)
        batch_size,sent_len,hidden_size = lstm_out.shape
        #Reshapes the output as a flat sequence compatible with softmax and loss functions
        lstm_out = lstm_out.contiguous().view(batch_size*sent_len,hidden_size)                    #lstm_out  [dim] = (batch_size*sent_len) x hidden_size
        return torch.tanh(lstm_out)

    def decode(self,tokens,pred_lexaction,pred_ytokens,pred_structaction,pred_structlabels,true_length):
        """
        This decodes a batched prediction into a tree. (~ performs best first parsing)
        Args:
           tokens            (tensor) : a list of integers (the true input tokens)
           pred_lexaction    (tensor) : a word sync tensor encoding lexical actions
           pred_ytokens      (tensor) : a word sync tensor encoding lexical labels
           pred_structaction (tensor) : a word sync tensor encoding structural actions
           pred_structlabels (tensor) : a word sync tensor encoding structural labels
           true_length          (int) : true length of the input (excludes padding)
        Returns.
           A couple. A parse derivation and a parse tree (or None if failure ?)
        """
        def decode_structural(nonterminal_label,nonterminal_action):
            """
            Converts back integer codes to strings for the structural case
            """
            action = self.ref_set.struct_action_vocab.itos[nonterminal_action]
            label  = self.ref_set.struct_vocab.itos[nonterminal_label]
            return label,action

        def decode_lexical(token,lexaction=None):
            """
            Converts back integer codes to strings for the lexical case
            """
            label  = self.ref_set.lex_vocab.itos[token]
            if lexaction:
                action = self.ref_set.lex_action_vocab.itos[lexaction]
                return label,action
            return label

        lextokens     = pred_ytokens.cpu().numpy() #use torch.no_grad before parsing code and remove the detach call
        lexactions    = pred_lexaction.cpu().numpy()
        structlabs    = pred_structlabels.cpu().numpy()
        structactions = pred_structaction.cpu().numpy()

        shift_init_c   = self.ref_set.lex_action_vocab.token_index( LCmodel.ACTION_SHIFT_INIT )
        shift_attach_c = self.ref_set.lex_action_vocab.token_index( LCmodel.ACTION_SHIFT_ATTACH )
        predict_c      = self.ref_set.struct_action_vocab.token_index( LCmodel.ACTION_PREDICT )
        attach_c       = self.ref_set.struct_action_vocab.token_index( LCmodel.ACTION_ATTACH )
        struct_pad_c   = self.ref_set.struct_action_vocab.token_index( self.ref_set.pad )
        struct_sos_c   = self.ref_set.struct_action_vocab.token_index( self.ref_set.sos )
        lex_pad_c      = self.ref_set.lex_action_vocab.token_index( self.ref_set.pad )

        Stack,Buffer  = [  ] , tokens.cpu().numpy()
        derivation    = [  ]
        logprob       = 0.0

        N = true_length
        r,d = 0,0
        for (idx,token,laction,ntlabel,ntaction) in zip(range(N),lextokens,lexactions,structlabs,structactions):
            #STRUCTURAL STATE
            d = len(Stack) if Stack else 0  #stack depth
            r = N - 1 - idx           #remaining words
            if idx > 0 : #skips start of sentence dummy state
                #preconditions (we mask for log probs)
                ntaction[ struct_pad_c ] = np.NINF
                ntaction[ struct_sos_c]  = np.NINF
                if d > 0 and not Stack[-1].is_complete():
                    ntaction[ predict_c ] = np.NINF
                if d > r+1:
                    ntaction[ predict_c ]  = np.NINF
                if d == 1:
                    ntaction[ attach_c ]  = np.NINF
                #decision
                #print('ntlabel',list(zip(self.ref_set.struct_vocab.itos,np.exp(ntlabel))))
                #print(idx,'struct_action',list(zip(self.ref_set.struct_action_vocab.itos,np.exp(ntaction))))
                ntlabel,struct_action = decode_structural(np.argmax(ntlabel),np.argmax(ntaction))
                #print('  ',struct_action)
                #exec
                if struct_action ==  LCmodel.ACTION_PREDICT :
                    Stack[-1] = LCtree(ntlabel,children=[Stack[-1]])
                elif struct_action == LCmodel.ACTION_ATTACH :
                    rc = Stack[-2].rightmost_attach_site()
                    rc.add_child(LCtree(ntlabel,children=[Stack[-1]]))
                    Stack.pop( )
                else:
                    print('structural action problem',struct_action)
                derivation.append( (struct_action,ntlabel,idx) )
            #LEXICAL STATE
            d = len(Stack) if Stack else 0  #stack depth
            #preconditions
            laction[ lex_pad_c ] = np.NINF
            if d == 0 :
                laction[ shift_attach_c ] = np.NINF
            if d > 0 and Stack[-1].is_complete():
                laction[ shift_init_c ] = np.NINF
            if d >= r+1:
                laction[ shift_init_c ] = np.NINF
            #decision
            ytoken,lex_action = decode_lexical(np.argmax(token), np.argmax(laction)) #pick the relevant prob for the token here ! (to be reworked)
            #print(idx,'lex_action',list(zip(self.ref_set.lex_action_vocab.itos,np.exp(laction))))
            #print(lex_action)

            #exec
            b0 = decode_lexical(Buffer[idx]) #forces the xtoken to be the reference rather than the predicted one
            if lex_action ==  LCmodel.ACTION_SHIFT_INIT:
                b0 = decode_lexical(Buffer[idx])
                Stack.append( LCtree(b0) )
            elif lex_action == LCmodel.ACTION_SHIFT_ATTACH :
                rc = Stack[-1].rightmost_attach_site()
                rc.add_child( LCtree(b0) )
            else:
                print('lexical action problem',struct_action)
            derivation.append( (lex_action,b0,idx) )

        #print("Derivation",derivation)
        #print("stack tree",Stack[-1])
        #print('derivation',derivation)
        if len(Stack) != 1:
            print("illegal stack",len(Stack))
        return derivation, Stack[-1]

    def eval_language_model(self,eval_set,batch_size=1,device=-1):
        """
        Only evaluates the core language model without parsing (only perplexity)
        Args:
           lm_set (ParsingDataSet): the development set
           batch_size        (int): the size of the batch
        Returns:
          float. The perplexity on this dev set
        """
        with torch.no_grad():
            dataloader = BucketLoader(eval_set,batch_size,device)
            N   = 0
            NLL = 0
            for batch in dataloader:

                seq_representation  = self.forward_base(batch.xtokens,batch.tokens_length)
                pred_ytokens,loss   = self.forward_lexical_tokens(seq_representation,batch.ytokens)

                locN = sum(batch.tokens_length)
                N   += locN
                NLL += (loss.item() * locN)  #lib implementation uses a mean rather than a basic sum
            return np.exp(NLL/N)

    def train_language_model(self,train_set,dev_set,epochs,batch_size=64,learning_rate=1.0,device=-1,alpha=0.0,clip=0.5,save_path='default'):
        """
        This trains a language model only (that can be used as a submodel of the parser).
        Meant to be used on very large data sets (such as the billion words corpus).
        Args :
          train_set (ParsingDataSet): xxx
          dev_set   (ParsingDataSet): xxx
          epochs               (int): xxx
        """
        #optimizer = optim.Adam(self.parameters(),lr=learning_rate)
        #optimizer = optim.ASGD(self.parameters(),lr=learning_rate,t0=5000,lambd=1,alpha=1)
        optimizer = optim.SGD(self.parameters(),lr=learning_rate)
        #scheduler = ReduceLROnPlateau(optimizer, mode='min', min_lr=0.001, factor=0.1, patience=10, verbose=True)
        scheduler = LambdaLR(optimizer,lr_lambda=lambda epoch:learning_rate/(1+epoch)**0.5)
        min_ppl   = 10000000000

        print('Starting...\n\n',file=sys.stderr,flush=True)

        for e in range(epochs):

            NLL = 0
            N   = 0
            dataloader = BucketLoader(train_set,batch_size,device,alpha)

            idx   = 0
            bsize = 0
            for batch in tqdm.tqdm(dataloader,total=dataloader.nbatches()):

                self.zero_grad( )
                seq_representation  = self.forward_base(batch.xtokens,batch.tokens_length)
                pred_ytokens,loss   = self.forward_lexical_tokens(seq_representation,batch.ytokens)
                loss.backward()
                clip_grad_norm_(self.parameters(), clip)
                optimizer.step()

                locN   = sum(batch.tokens_length)
                NLL   += (loss.item() * locN)
                N     += locN
                bsize += len(batch.tokens_length)
                if idx > 0 and idx % 1000 == 0:
                    ppl = self.eval_language_model(dev_set,batch_size,device)
                    print('PPL',ppl,'mean batch size',bsize/idx,'LR =',optimizer.param_groups[0]['lr'],file=sys.stderr,flush=True)
                    #scheduler.step(ppl)
                    scheduler.step()
                idx +=1

            print("\nEpoch",e,'training loss (NLL) =', NLL/N ,'training PPL =',np.exp(NLL/N), 'learning rate =',optimizer.param_groups[0]['lr'],flush=True)
            ppl = self.eval_language_model(dev_set,batch_size,device)
            print("                                                dev      PPL =",ppl,flush=True)
            if ppl < min_ppl:
                min_ppl = ppl
                print('     => model saved.',flush=True)
                self.save(save_path)

    def eval_parser(self,dev_set,batch_size=1,device=-1,with_loss=False):
        """
        Evaluates the parser on a dev set.
        Args:
           dev_set (ParsingDataSet): the development set
           batch_size         (int): the size of the batch
        Returns:
           a list of trees. Returns the predicted data set as whole. The original ordering of dev_set is guaranteed to be preserved.
        """
        with torch.no_grad():

            dataloader = BucketLoader(dev_set,batch_size,device)
            orig_idxes = [ ]
            pred_trees = [ ]

            N   = 0
            NLL = 0

            if with_loss:
                lex_action_loss    = nn.NLLLoss(reduction='sum',ignore_index=train_set.lex_action_vocab.stoi[train_set.pad])
                struct_action_loss = nn.NLLLoss(reduction='sum',ignore_index=train_set.struct_action_vocab.stoi[train_set.pad])
                struct_loss        = nn.NLLLoss(reduction='sum',ignore_index=train_set.struct_vocab.stoi[train_set.pad])

            for batch in dataloader:
                seq_representation =  self.forward_base(batch.xtokens,batch.tokens_length)

                pred_lexaction,loss1     =  self.forward_lexical_actions(seq_representation,ref_action=batch.lex_actions,loss=lex_action_loss)
                pred_structaction,loss2  =  self.forward_structural_actions(seq_representation,ref_action=batch.struct_actions,loss=struct_action_loss)
                pred_ytokens,loss3       =  self.forward_lexical_tokens(seq_representation,batch.ytokens)
                pred_structlabels,loss4  =  self.forward_structural_labels(seq_representation,ref_action=batch.struct_labels,loss=struct_loss)

                if with_loss:
                    NLL += loss1.item() + loss2.item() + loss3.item() + loss4.item()
                    N   += sum(batch.tokens_length)

                #here we reshape sentence_wise :
                #   input  [dim] = (batch_size*sent_len) x hidden_size
                batch_size,batch_len = batch.ytokens.shape
                #print('eval batch size x batch len (', batch_size,batch_len,')')
                #   output [dim] = batch_size x sent_len x hidden_size

                pred_lexaction    = pred_lexaction.view(batch_size,batch_len,-1)
                pred_structaction = pred_structaction.view(batch_size,batch_len,-1)
                pred_ytokens      = pred_ytokens.view(batch_size,batch_len,-1)
                pred_structlabels = pred_structlabels.view(batch_size,batch_len,-1)

                orig_idxes.extend(batch.orig_idxes)
                batch_preds = [ self.decode(batch.ytokens[idx],\
                                            pred_lexaction[idx],\
                                            pred_ytokens[idx],\
                                            pred_structaction[idx],\
                                            pred_structlabels[idx],\
                                            batch.tokens_length[idx]) for idx in range(batch_size)]

                pred_trees.extend(batch_preds)

            matched_idxes = enumerate(orig_idxes) #iterates using the ascending original order of the data set
            if with_loss:
                print("        development loss   (NLL) = ", NLL/(4*N),flush=True)
            return [ pred_trees[current_idx] for (current_idx,orig_idx) in sorted(matched_idxes,key=lambda x:x[1]) ]

    def train_parser(self,train_set,dev_set,epochs,batch_size=1,learning_rate=0.1,device=-1,alpha=0.0,save_path='.'):
        """
        Args :
          train_set (ParsingDataSet): xxx
          dev_set   (ParsingDataSet): xxx
          epochs               (int): xxx
        """
        lex_action_loss    = nn.NLLLoss(reduction='sum',ignore_index=train_set.lex_action_vocab.stoi[train_set.pad])
        struct_action_loss = nn.NLLLoss(reduction='sum',ignore_index=train_set.struct_action_vocab.stoi[train_set.pad])
        struct_loss        = nn.NLLLoss(reduction='sum',ignore_index=train_set.struct_vocab.stoi[train_set.pad])
        optimizer          = optim.SGD(self.parameters(),lr=learning_rate)
        scheduler          = ReduceLROnPlateau(optimizer, mode='min', factor=0.1 ,patience=10, verbose=True)

        for e in range(epochs):

            _lex_loss,_lex_action_loss, _struct_action_loss, _struct_loss = 0,0,0,0
            N = 0

            dataloader = BucketLoader(train_set,batch_size,device,alpha)

            for batch in dataloader:

                self.zero_grad()

                seq_representation =  self.forward_base(batch.xtokens,batch.tokens_length)

                pred_lexaction,loss1     =  self.forward_lexical_actions(seq_representation,ref_action=batch.lex_actions,loss=lex_action_loss)
                pred_structaction,loss2  =  self.forward_structural_actions(seq_representation,ref_action=batch.struct_actions,loss=struct_action_loss)
                pred_ytokens,loss3       =  self.forward_lexical_tokens(seq_representation,batch.ytokens)
                pred_structlabels,loss4  =  self.forward_structural_labels(seq_representation,ref_action=batch.struct_labels,loss=struct_loss)

                _lex_loss           += loss3.item()
                _lex_action_loss    += loss1.item()
                _struct_loss        += loss4.item()
                _struct_action_loss += loss2.item()
                N += sum(batch.tokens_length)

                loss1.backward(retain_graph=True)
                loss2.backward(retain_graph=True)
                loss3.backward(retain_graph=True)
                loss4.backward()
                optimizer.step()

            L = _lex_loss + _lex_action_loss + _struct_action_loss + _struct_loss
            print("Epoch",e,'training loss (NLL) =', L/(4*N),'learning rate =',optimizer.param_groups[0]['lr'],flush=True)#,scheduler.get_lr()[0],N)
            print('        lex loss           (NLL) = ',_lex_loss/N,flush=True)
            print('        lex action loss    (NLL) = ',_lex_action_loss/N,flush=True)
            print('        struct loss        (NLL) = ', _struct_loss/N,flush=True)
            print('        struct action loss (NLL) = ',_struct_action_loss/N,flush=True)
            scheduler.step(L)
            #Development f-score computation
            pred_trees = list(tree for (derivation,tree) in self.eval_parser(dev_set,batch_size,device,with_loss=True))
            #for t in pred_trees[:10]:
            #    print(t)
            fscores    = [ reftree.compare(predtree)[2]   for (predtree,reftree) in zip(pred_trees,dev_set.tree_set) ]
            print("        development F-score      = ", sum(fscores) / len(fscores),flush=True)
    @staticmethod
    def derivation2tree(derivation):
        """
        This builds up a ConsTree from a derivation
        Args:
           derivation (list):a list of actions
        Returns:
           ConsTree. a constituent tree object built from the derivation
        """
        derivation.reverse()
        tree_stack = [ ]
        while derivation:
            action,label = derivation.pop()
            if action == LCmodel.ACTION_SHIFT_INIT :
                tree_stack.append( ConsTree(label) )
            elif action == LCmodel.ACTION_PREDICT:
                left_child = tree_stack.pop()
                tree = ConsTree(label,children=[ left_child ])
                tree_stack.append(tree)
            elif action == LCmodel.ACTION_SHIFT_ATTACH:
                LCmodel.rightmostNT(tree_stack[-1]).add_child(ConsTree(label))
            elif action == LCmodel.ACTION_ATTACH:
                child = tree_stack.pop()
                child = ConsTree(label,children=[child])
                LCmodel.rightmostNT(tree_stack[-1]).add_child(child)
            else:
                print('derivation error')
        return tree_stack[-1]

    @staticmethod
    def oracle_derivation(ctree,left_corner=True):
        """
        This computes an oracle derivation from a constituent tree.
        The tree must be strictly binarized. (Only binary rules and
        nothing else)
        Args:
           ctree   (ConsTree): the tree to extract the derivation from
           left_corner (bool): the tree is a current left corner
           root        (bool): the tree is the topmost node in the overall tree
        Returns:
           a list of actions.
        """
        if ctree.is_leaf( ):
            return [(LCmodel.ACTION_SHIFT_INIT,ctree.label)] if left_corner else [(LCmodel.ACTION_SHIFT_ATTACH,ctree.label)]
        else: #internal node
            derivation = LCmodel.oracle_derivation(ctree.get_child(0),left_corner=True)
            if left_corner:
                derivation.append((LCmodel.ACTION_PREDICT,ctree.label))
            else:
                derivation.append((LCmodel.ACTION_ATTACH,ctree.label))
            derivation.extend( LCmodel.oracle_derivation(ctree.get_child(1),left_corner=False) )
        return derivation

def input_treebank(filename):
    """
    Data generator that reads and preprocesses a treebank from a file
    Args:
       filename (string): a file where to read trees
    Yields:
       a sequence of trees
    """
    istream = open(filename)
    for treeline in istream:
        tree = ConsTree.read_tree(treeline)
        tree.add_eos()
        ConsTree.right_markovize(tree)
        ConsTree.close_unaries(tree)
        tree.strip_tags()
        yield tree
    istream.close()

def output_treebank(treelist,filename=None):
    """
    Prints a treelist to file
    Args:
       treelist   (list): a list of trees
       filename (string): a file where to write trees
    """
    ostream = sys.stdout if filename is None else open(filename,'w')
    for tree in treelist:
        tree.expand_unaries()
        tree.unbinarize( )
        tree.strip_eos( )
        print(tree,file=ostream)
    if filename:
        ostream.close()

if __name__ == '__main__':

    import argparse
    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter,description='Trains and predicts parsing language models on GPUs with a left corner algorithm. MODELDIR and VOCABNAME are mandatory arguments. The parser performs training if both TRAIN and DEV are provided and testing if TEST is provided. A new language model is trained if TRAIN, DEV and LMTRAIN are provided. A new vocabulary is generated if VOCABTRAIN is provided too. If MODELDIR is already existing, this existing model is updated. Training both a language model and a parsing model requires to run the command twice (with update).')
    parser.add_argument('--epochs', dest='epochs',         type=int,default=10,help='max number of epochs')
    parser.add_argument('--lr', dest='lr',                 type=float,default=0.001,help='Learning rate of the trainer')
    parser.add_argument('--rnn-memory', dest='rnn_memory', type=int,default=600,help='Size of the RNN memory')
    parser.add_argument('--device', dest='device',       type=int,default=-1,help='Device (GPU) number where to run the computations: -1 for CPU, 0 or more for GPUs')
    parser.add_argument('--embedding', dest='embedding', type=int,default=300,help='Size of the embeddings')
    parser.add_argument('--batch', dest='batch_size', type=int,default=32,help='Size of the mini-batches')
    parser.add_argument('--lm-train', dest='lmtrain', type=str,help='Path to the root of the billion word dataset')
    parser.add_argument('--vocab-train', dest='vocabtrain', type=str,help='Path to the corpus where to extract the vocab')
    parser.add_argument('--train', dest='train', type=str,help='Path to the train treebank')
    parser.add_argument('--dev', dest='dev',     type=str,help='Path to the development treebank')
    parser.add_argument('--test', dest='test',     type=str,help='Path to the test treebank')
    parser.add_argument('--model-dir', dest='modeldir',     type=str,help='Path to the model directory',required=True)
    parser.add_argument('--vocab-threshold', dest='lex_min_freq', type=int,default=100,help='Builds a new vocabulary from VOCABTRAIN dataset with words of frequency >= threshold.')
    parser.add_argument('--vocab-name', dest='vocabname', type=str,help='Path to the lexical vocabulary file',required=True) #e.g. vocab100

    args = parser.parse_args()

    if args.vocabtrain and args.lex_min_freq:
        print('Extracting vocabulary from dir %s with words of frequency >= %d'%(args.vocabtrain,args.lex_min_freq),file=sys.stderr,flush=True)
        evocab = Vocabulary(extract_vocabulary(args.vocabtrain),unk='<unk>',sos='<sos>',min_freq=args.lex_min_freq) #'/home/bcrabbe/parsing_as_LM/rnng/history/billion_words'; 100
        evocab.save( args.vocabname )
        print('Vocab saved as %s'%(args.vocabname,),file=sys.stderr,flush=True)
    else:
        evocab = Vocabulary.load( args.vocabname )

    if args.train and args.dev:
        print('Loading treebank data...',file=sys.stderr)
        trainset = list(input_treebank(args.train))
        devset   = list(input_treebank(args.dev))
        train_df = ParsingDataSet(trainset,ext_vocab=evocab)
        dev_df   = ParsingDataSet(devset,root_dataset=train_df)
        print('treebank data loaded.',file=sys.stderr,flush=True)
        if args.lmtrain :
            print('Loading language model data...',file=sys.stderr)
            lmset    = list(load_billion_full(args.lmtrain))  #'/home/bcrabbe/parsing_as_LM/rnng/history/billion_words'
            lm_df    = ParsingDataSet(lmset,ext_vocab=evocab)
            print('\nlanguage model data loaded.',file=sys.stderr,flush=True)

        if os.path.exists(args.modeldir):
            print('Existing model detected at %s. Training will update this model.'%(args.modeldir,),file=sys.stderr,flush=True)
            parser = LCmodel.load(args.modeldir,device=args.device)
            update_suff = '.up'
        else:
            print('New model will be saved as %s'%(args.modeldir,),file=sys.stderr,flush=True)
            parser = LCmodel(train_df,rnn_memory=args.rnn_memory,embedding_size=args.embedding,device=args.device)
            update_suff = ''

        if args.lmtrain:
            parser.train_language_model(lm_df,dev_df,args.epochs,batch_size=args.batch_size,learning_rate=args.lr,device=args.device,alpha=0.0,save_path=args.modeldir+update_suff)
        else:
            if not update_suff:
                print('Cannot train parser from scratch. Train language model first.\naborting.',file=sys.stderr)
                exit(1)
            parser.train_parser(train_df,dev_df,args.epochs,batch_size=args.batch_size,learning_rate=args.lr,device=args.device,alpha=0.0,save_path=args.modeldir+update_suff)

        print('Done. Model saved as %s'%(args.modeldir+update_suff),file=sys.stderr,flush=True)

    if args.test:
        print('Running test...',file=sys.stderr,flush=True)
        parser = LCmodel.load(args.modeldir,device=args.device)
        #Todo (load a list of strings and generate a test set with help from the parser)
        #parser.eval_parser(dev_set,batch_size=1,device=-1,with_loss=False):
        #Todo generate trees and unpack them


def breadth_first(active_set, s1, a, mg):
    active_set.append((0, s1))
    return active_set

class GroundingSpace:
    """
    fork of grounding space from esdcs
    """
    def __init__(self, context=None, start_state=None, event_depth=1, salient_tags=[]):
        self.start_state = start_state
        self.context = context
        self.event_depth = event_depth
        self.objects = [o for o in context.objects]
        # self.objects = [o for o in context.objects if o != context.agent]
        self.salient_tags = salient_tags

        if self.start_state == None:
            raise ValueError("Must pass in a start state.")

        assert self.context != None
        assert self.start_state != None

        self.context = context
        self.events = self.events_as_search()
        #self.paths = [p.path for events in self.events for p in events]
        #assert len(self.events) > 0
        #assert len(self.paths) > 0

    def events_as_search(self, callback=breadth_first):
        movable_objects = [o for o in self.context.objects if "pallet" in o.tags]
        pobj_lists = []
        states = []

        active_set = [(0, self.start_state)]

        for i in range(self.event_depth):
            new_active_set = []
            for cost, s in active_set:
                for (s1, a, mg) in s.getSuccessors(groundings=movable_objects):
                    pobj_lists.append(mg)
                    new_active_set = callback(new_active_set, s1, a, mg)
                    states.append((s1, a, mg))

            active_set = new_active_set
        return states

    def grounding_space(self, node):
        """
        Returns the space of possible values for a node, given the context.
        For a node which is already bound to a value (as in the case of yes/no dialog), returns only the value for that node.
        """
        if node.is_object:
            if len(self.salient_tags) > 0:
                return [o for o in self.objects if not set(self.salient_tags).isdisjoint(set(o.tags))]
            return self.objects
        elif node.is_place:
            return self.context.places
        elif node.is_event:
            return self.events
        else:
            raise ValueError("Unexpected node type: %s" % node)


from datetime import datetime
from pathlib import Path

from .utils import split_into_parts
import schedlib as sl
from schedlib.policies import BasicPolicy
import random

random.seed(int(datetime.now().timestamp()))
default_schedule = Path(__file__).parent / "schedule_sat.txt"

def dummy_policy(t0, t1, policy_config={}, app_config={}):
    dt = abs(t1.timestamp() - t0.timestamp())  # too lazy to check for t1<t0 now
    # get current time as a timestamp string
    now = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    n = random.randint(1, 10)
    parts = split_into_parts(dt, n)
    commands = ["import time", f"# {now}"]
    for i, part in enumerate(parts):
        commands += [f"time.sleep({part:.2f})"]
    commands = "\n".join(commands)
    return commands

def basic_policy(t0, t1, policy_config, app_config={}):
    policy = BasicPolicy(**policy_config)
    seq = policy.init_seqs(t0, t1)
    seq = policy.apply(seq)
    cmd = policy.seq2cmd(seq)
    return str(cmd)


#coding:utf-8

import requests
import read

#获取参保日志
def ss_getSsOpFlowLogs_get():
    body = {
        "userId":read.userId(),
        "flag":"true"   #true:查所有  false:查最新3条
    }

    url = read.url_post() + 'getSsOpFlowLogs'
    r = requests.get(url,json=body,headers=read.header())

    #检测返回值
    print(r.text)
    #检测返回状态
    print(r.status_code)

if __name__ == '__main__':
    ss_getSsOpFlowLogs_get()


a = []
b = 0
while True:
    b = int(input('Digite um número qualquer: '))
    if not a or a[0] < b:
        a.insert(0, b)
    else:
        for i in range(len(a)-1, -1, -1):
            if b < a[i]:
                a.insert(i+1, b)
                break
    b = ''
    while True:
        b = input('Quer continuar? [S/N] ')
        if b in 'Nn':
            break
        elif b in 'Ss':
            break
        else:
            print('Você digitou errado tente novamente')
    if b in 'Nn':
        break
print(f'foram digitados {len(a)} vezes, e a lista é: {a}', end='')
if a.count(5) == 1:
    print(f' o 5 está na lista na posição {a.index(5)}')
else:
    print(' não a número(s) 5')


import pandas as pd
from progressbar import *
import numpy as np

def remove_duplicate(df, has):
    cnt = 0
    widgets = ['Progress: ', Percentage(), ' ', Bar('#'), ' ', Timer(),
               ' ', ETA()]
    pbar = ProgressBar(widgets=widgets, maxval=len(df)).start()
    res_df = df[0:1]
    has.add(df[0:1]["created_at"].values[0])
    for i in df.index:
        cur = df[i:i+1]
        if cur["created_at"].values[0] in has:
            cnt += 1
        else:
            has.add(cur["created_at"].values[0])
            res_df = res_df.append(cur)
        pbar.update(i+1)
    pbar.finish()
    print("Done! drop {} duplicate.".format(cnt))
    return res_df

has = set()
trump_df = pd.read_csv('Data/hashtag_donaldtrump.csv', lineterminator='\n')
print("Read Trump Data Complete!")
print("Start Processing Trump Data...")
trump_df["tag_biden"] = np.zeros(len(trump_df), dtype=np.int)
trump_df["tag_trump"] = np.ones(len(trump_df), dtype=np.int)
trump_df = trump_df.drop(['tweet', 'user_description'], axis=1)
res_df = remove_duplicate(trump_df, has)
del trump_df

biden_df = pd.read_csv('Data/hashtag_joebiden.csv', lineterminator='\n')
print("Read Trump Data Complete!")
print("Start Processing Biden Data...")
biden_df["tag_biden"] = np.ones(len(biden_df), dtype=np.int)
biden_df["tag_trump"] = np.zeros(len(biden_df), dtype=np.int)
biden_df = biden_df.drop(['tweet', 'user_description'], axis=1)
res_df.append(remove_duplicate(biden_df, has))

res_df.to_csv("election_twitter_data.csv", index=False)
# biden_df = pd.read_csv('pycode\hashtag_joebiden.csv',lineterminator='\n')


import numpy as np
import torch
import time
from sko.GA import GA
import flopy.utils.binaryfile as bf
import pandas as pd
import flopy
import matplotlib.pyplot as plt

def MOD(a,b,c,d):            # MODFLOW程序
    model = flopy.modflow.Modflow.load(r"E:\文件资料\大论文\MODFLOW\genetic\model_MODFLOW_text\model.mfn",
                                            exe_name = r'D:\GMS10.5\models\mf2k5\usgs\mf2005.exe')
    wel_spd = {0: [[0, 0, 0, a], [0, 0, 9,b],[0,6,0,c],[0,6,9,d]],      #[lay, row, col, flux]
              1: [[0, 0, 0, a], [0, 0,9,b],[0,6,0,c],[0,6,9,d]]}
    wel = flopy.modflow.ModflowWel(model, stress_period_data = wel_spd)
    model.write_input()
    model.run_model()
    hds = bf.HeadFile('model' + ".hed") # 获取水头文件
    head = hds.get_data() # 提取水头值
    df1 = pd.DataFrame(head[0])
    plt.subplot(121)
    plt.title('head')
    plt.imshow(df1,cmap='RdBu')
    plt.colorbar()
    plt.show(block=False)

    # df = df.drop([0,1,2,7,8,9],axis=0)       # 删除多余行
    # df = df.drop([0,1,2,7,8,9],axis=1)       # 删除多余列
    水位值 = df1.iloc[4,4]
    print(水位值)

    vds = bf.HeadFile('model' + ".vert_dis",text='Z DISPLACEMENT') # 获取沉降文件
    vd = vds.get_data() # 提取沉降值
    df2 = pd.DataFrame(vd[0])
    plt.subplot(122)
    plt.title('subsidence')
    plt.imshow(df2,cmap='Greys')
    plt.colorbar()
    plt.pause(0.001)
    plt.clf()

    沉降值 = df2.iloc[4,4]
    print(沉降值)

    return 水位值,沉降值

def obj_func(p):
    a,b,c,d,e,f,g,h= p
    水位值,沉降值 = MOD(a*e,b*f,c*g,d*h)
    A = e*(120 + 1.2*np.abs(a))
    B = f*(130 + 1.3*np.abs(b))
    C = g*(140 + 1.4*np.abs(c))
    D = h*(150 + 1.5*np.abs(d))
    return 100*np.abs(水位值-98)+100*np.abs(沉降值)+A+B+C+D

fig=plt.figure(figsize=(9,3))

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
ga = GA(func=obj_func, n_dim=8, size_pop=100, prob_mut=0.01,max_iter=10, lb=[-50,-50,-50,-50,0,0,0,0], ub=[0,0,0,0,1,1,1,1],
        precision=[1e-3,1e-3,1e-3,1e-3,1,2,3,4])
ga.to(device=device)
start_time = time.time()
best_x, best_y = ga.run()
print(time.time() - start_time)
print('best_x:', best_x, '\n', 'best_y:', best_y)

Y_history = pd.DataFrame(ga.all_history_Y)
fig, ax = plt.subplots(1, 1)
Y_history.min(axis=1).cummin().plot(kind='line')
plt.show()


# PyBank Plan of Attack:
    # Resource data
    #   Col1: Date | MMM-YY | Min: 1/1/2010 | Max: 2/1/2017 |
    #   Col2: Profit/Losses | +/- # | Max: 1170593 | Min: -1196225 |

    # Deliverables:
    #   1.  The total number of months included in the dataset
    #   2.  The net total amount of "Profit/Losses" over the entire period
    #   3.  The average of the changes in "Profit/Losses" over the entire period
    #   4.  The greatest increase in profits (date and amount) over the entire period
    #   5.  The greatest decrease in losses (date and amount) over the entire period
    #   6.  Write the results to a text file.

# HW Attempt: PyBank
# Import modules
import os
import csv

# Set the path
datapath = os.path.join('..','Resources', 'PyBank_Data.csv')

# Read the csv data
with open(datapath) as csvfile:
    reader = csv.reader(csvfile, delimiter=",")
    # Store header row
    headers = next(reader)

    # Create separate lists for the date and P/Ls field
    dates = []
    profits_losses = []

    for row in reader:
        dates.append((row[0]))
        profits_losses.append((row[1]))

    # Stores the P&L list as integers
    profits_losses = [int(i) for i in profits_losses]

# 1. Number of entries
    no_months = len(dates)

# 2. Loop to get the net profit from all entries
    net_profit = 0
    for x in profits_losses:
        net_profit += x

# 3. Average Monthly Change
    # (referenced stackoverflow: https://stackoverflow.com/questions/2400840/finding-differences-between-elements-of-a-list)
    monthly_change = [j - i for i , j in zip(profits_losses[:-1],profits_losses[1:])]

    # Total # of changes
    num_changes = len(monthly_change)

    #loop for find the sum of all changes
    sum_change = 0
    for x in monthly_change:
        sum_change += x

    avg_change = (sum_change / num_changes)
    rd_avg = round(avg_change,2)

# 4. Find the max profit gain in a month
    g_inc = max(profits_losses)
    index_g_inc = profits_losses.index(g_inc)
    date_g_inc = dates[index_g_inc]

# 5. Find the min and max profits
    g_dec = min(profits_losses)
    index_g_dec = profits_losses.index(g_dec)
    date_g_dec = dates[index_g_dec]

# 6. Print the results and write to txt file
    print("----------------------------")
    print("Financial Analysis:")
    print("")
    print(f"Total # of Months: {no_months}")
    print(f"Net Profit: ${net_profit}")
    print(f"Average Monthly Change: ${rd_avg}")
    print(f"Greatest Profit Increase: ${g_inc} in {date_g_inc}")
    print(f"Greatest Profit Decrease: ${g_dec} in {date_g_dec}")
    print ("----------------------------")

    resultspath = os.path.join('..','Analysis', 'Results_pybank.txt')
    with open(resultspath , "w") as text:
        text.write("----------------------------")
        text.write("\nFinancial Analysis:")
        text.write("\n")
        text.write(f"\nTotal # of Months: {no_months}")
        text.write(f"\nNet Profit: ${net_profit}")
        text.write(f"\nAverage Monthly Change: ${rd_avg}")
        text.write(f"\nGreatest Profit Increase: ${g_inc} in {date_g_inc}")
        text.write(f"\nGreatest Profit Decrease: ${g_dec} in {date_g_dec}")
        text.write ("\n----------------------------")


#!/usr/bin/env python3
# coding=utf-8

from random import randint
import matplotlib.pyplot as plt

# Oppgave 1.a
def kurvediagram() -> None:
    """
    Plotter kurver av 10 terningkast som viser utviklingen av:
        gjennomsnittsverdien oppnådd med terninger så langt
        minimumsverdien oppnådd med terninger så langt
        maksimumsverdien oppnådd med terninger så langt.
    """
    sum = 0
    minimum = [6]
    maximum = [0]
    gjennomsnitt= []
    for i in range(0, 10):
        kast = randint(1, 6)
        sum += kast
        if i != 0:
            gjennomsnitt.append(sum / i)
        else:
            gjennomsnitt.append(kast)
        if kast >= maximum[-1]:
            maximum.append(kast)
        else:
            maximum.append(maximum[-1])
        if kast <= minimum[-1]:
            minimum.append(kast)
        else:
            minimum.append(minimum[-1])
    minimum.pop(0)
    maximum.pop(0)
    x = range(1, 11)
    plt.plot(x, gjennomsnitt, 'r', label='Gjennomsnitt')
    plt.plot(x, minimum, 'g', label='Minste')
    plt.plot(x, maximum, 'b', label='Største')
    plt.legend()
    plt.show()

# Oppgave 1.b
def punktdiagram() -> None:
    """
    Lager et punktdiagram av 10 kast med 2 terninger
    """
    (terning1, terning2) = ([], [])
    for _ in range(0, 10):
        terning1.append(randint(1,6))
        terning2.append(randint(1,6))
    plt.scatter(terning1, terning2)
    plt.show()

# Oppgave 1.c
def stolpediagram() -> None:
    """
    Lager et stolpediagram av hvor mange av 100 terningkast
    som ga verdiene 1 - 6
    """
    x = range(1, 7)
    y = [0, 0, 0, 0, 0, 0]
    for _ in range(0, 100):
        kast = randint(1, 6)
        y[kast - 1] += 1
    plt.bar(x, y)
    plt.show()

def dobbel_stolpediagram() -> None:
    """
    Lager et stolpediagram av hvor mange av 100 terningkast
    med 2 terninger(rød og blå) som ga verdiene 1 - 6
    """
    x = range(1, 7)
    y_blå = [0]*6
    y_rød = [0]*6
    for _ in range(0, 100):
        blå_kast = randint(1, 6)
        rød_kast = randint(1, 6)
        y_blå[blå_kast - 1] += 1
        y_rød[rød_kast - 1] += 1
    plt.bar(x, y_blå, width=-0.4, align='edge', color='blue')
    plt.bar(x, y_rød, width=0.4, align='edge', color='red')
    plt.show()

if __name__ == "__main__":
    kurvediagram()
    punktdiagram()
    stolpediagram()
    dobbel_stolpediagram()


'''
40に加えて、分節を示すクラスChunkの実装
形態素（Morphオブジェクト）のリスト（morphs），
係り先文節インデックス番号（dst），係り元文節インデックス番号のリスト（srcs）をメンバ変数
冒頭の説明文の文節の文字列と係り先を表示
'''

# 文節
class Chunk:
  def __init__(self,morphs,dst):
    # 形態素のリスト(morphs)
    self.morphs = morphs
    # 係り先
    self.dst = dst
    # 係り元文節の保存
    self.srcs = []

class Morph:
  def __init__(self,block):
    surface,attr = block.split('\t')
    attr = attr.split(',')
    self.surface = surface
    self.base = attr[6]
    self.pos = attr[0]
    self.pos1 = attr[1]

# 文節(chunks)リスト
class Sentence:
  def __init__(self,chunks):
    self.chunks = chunks
    # i が文節番号
    # そもそもEOSで文節ごとに分割しているため、
    # その文節を超えた i は指定されることはない。
    for i ,chunk in enumerate(self.chunks):
      # 係先があるか
      if chunk.dst not in [None,-1]:
        # ex)i番目に17Dがあれば、17のchunk.srcsにiを格納
        self.chunks[chunk.dst].srcs.append(i)

filename = 'ai.ja/ai.ja.txt.parsed'
sentences = []
# chunkのインスタンス変数の集合
chunks = []
morphs = []

with open(filename) as f:
  for block in f:
    # Add Chunk
    if block[0] == "*":
      if len(morphs) > 0:
        chunks.append(Chunk(morphs,dst))
        morphs = []
      # 文字列の末尾部分を除去
      # 引数より除去される文字集合
      dst = int(block.split(' ')[2].rstrip('D'))
    # Add Morph
    elif block != 'EOS\n':
      morphs.append(Morph(block))

    # paragraph
    else:
      chunks.append(Chunk(morphs,dst))
      sentences.append(Sentence(chunks))
      morphs = []
      chunks = []
      dst = None # 初期化

for chunk in sentences[2].chunks:
  print([morph.surface for morph in chunk.morphs], chunk.dst, chunk.srcs)
'''
['人工', '知能'] 17 []
['（', 'じん', 'こうち', 'のう', '、', '、'] 17 []
'''