import csv
import os
import sys

import music21
import numpy as np
import xgboost

DeepGRU = 0# from approach_attention import DeepGRU
from loader_representations import get_path, velocity_piece, notes_piece, finger_piece, finger_nakamura_piece, \
    prob_piece
import torch
from torch.nn.utils.rnn import pack_padded_sequence as packer, pad_packed_sequence as padder

from utils import strm2map, load_json, save_json

def salamizer(X, win_size=5):
    X = np.array(X)
    X_ans = []
    hop_size = 1

    for i in range(1, X.shape[0] - win_size + 1, hop_size):
        window = X[i:i + win_size, :].reshape((-1, 1))  # each individual window
        X_ans.append(window)
    X_ans = np.squeeze(np.array(X_ans))
    return X_ans

def get_onset_difficulty(prediction, onsets):
    ans = {}
    for idx, p in enumerate(prediction):
        ans[onsets[idx]] = p
    return ans

def save_PIG_difficulty(alias, model, piece, onset_difficulty, rep):
    path_alias = get_path(alias)
    path_to_save = os.path.join('visualization', model, piece + '.txt')
    r_h_cost = '/'.join(["Fingers", path_alias, piece + '_rh.txt'])
    l_h_cost = '/'.join(["Fingers", path_alias, piece + '_lh.txt'])
    with open(r_h_cost) as tsv_file:
        r_h = list(csv.reader(tsv_file, delimiter="\t"))
    with open(l_h_cost) as tsv_file:
        l_h = list(csv.reader(tsv_file, delimiter="\t"))

    PIG_content = []
    for idx, content in enumerate(sorted(r_h + l_h, key=lambda a: float(a[1]))):
        new_content = content
        if content[7] != -1 and round(float(content[1]), 2) in onset_difficulty:
            new_content[0] = idx
            new_content.append(round(onset_difficulty[round(float(content[1]), 2)]))
            PIG_content.append(new_content)
        else:
            new_content[0] = idx
            new_content.append(-1)
            PIG_content.append(new_content)

    with open(path_to_save, 'w') as tsvfile:
        writer = csv.writer(tsvfile, delimiter='\t')
        for record in PIG_content:
            writer.writerow(record)

KEY_TO_SEMITONE = {'c': 0, 'c#': 1, 'db': 1, 'd-': 1, 'c##': 2, 'd': 2, 'e--': 2, 'd#': 3, 'eb': 3, 'e-': 3, 'd##': 4, 'e': 4, 'f-': 4, 'e#': 5,
                   'f': 5, 'g--': 5, 'e##': 6, 'f#': 6, 'gb': 6, 'g-': 6, 'f##': 7, 'g': 7, 'a--': 7, 'g#': 8, 'ab': 8, 'a-': 8,
                   'g##': 9, 'a': 9, 'b--': 9,
                   'a#': 10, 'bb': 10, 'b-': 10, 'a##': 11, 'b': 11, 'b#': 12, 'c-': -1, 'x': None}

def an2midi(an):
    a = an[:-1].lower()  # alpha
    n = int(an[-1])  # numeric
    return n * 12 + KEY_TO_SEMITONE[a]

def save_score_difficulty(alias, output, piece, onset_difficulty, rep, appr):

    is_nakamura = rep in ['prob', 'finger_nakamura']
    path_to_save = os.path.join(output, os.path.basename(piece)[:-4] + '.pdf')
    if not is_nakamura:
        r_h_cost = '/'.join(["Fingers", "pianoplayer", os.path.basename(piece)[:-4] + '_rh.txt'])
        l_h_cost = '/'.join(["Fingers", "pianoplayer", os.path.basename(piece)[:-4] + '_lh.txt'])
        with open(r_h_cost) as tsv_file:
            r_h = list(csv.reader(tsv_file, delimiter="\t"))
        with open(l_h_cost) as tsv_file:
            l_h = list(csv.reader(tsv_file, delimiter="\t"))
        h = sorted(r_h + l_h, key=lambda a: float(a[3]), reverse=True)
    else:
        h_cost = '/'.join(["Fingers", "nakamura", os.path.basename(os.path.basename(piece)[:-4]) + '.txt'])
        with open(h_cost) as tsv_file:
            all_h = list(csv.reader(tsv_file, delimiter="\t"))[1:]
        for idx in range(len(all_h)):
            all_h[idx][3] = an2midi(all_h[idx][3])
        h = sorted(all_h, key=lambda a: float(a[3]), reverse=True)
    r_diff = []
    l_diff = []

    for idx, content in enumerate(sorted(h, key=lambda a: float(a[1]))):
        if content[6] == '0':
            if (is_nakamura or content[7] != -1) and round(float(content[1]), 2) in onset_difficulty:
                r_diff.append((content[7], (round(onset_difficulty[round(float(content[1]), 2)]))))
            else:
                r_diff.append((content[7], -1))
        else:
            if (is_nakamura or content[7] != -1) and round(float(content[1]), 2) in onset_difficulty:
                l_diff.append((content[7], round(onset_difficulty[round(float(content[1]), 2)])))
            else:
                l_diff.append((content[7], -1))
    if appr == 'deepgru':
        INTERP = [
            '#000061',
            '#0000cc',
            '#0000ff',
            '#3333ff',
            '#6666ff',
            '#9999ff',
            '#b3b3ff',
            '#ccccff',
            '#e6e6ff',
            'white'
        ]
    else:
        green = '#a1de00'
        yellow = '#f6b100'
        red = '#e30000'
        INTERP = [green, yellow, red, 'white']
    sf = music21.converter.parse(piece)
    rh_om = strm2map(sf.parts[0])
    lh_om = strm2map(sf.parts[1])
    for om, diff_list in zip([rh_om, lh_om], [r_diff, l_diff]):
        for o, (finger, diff) in zip(om, diff_list):
            if 'chord' in o:
                music21_structure = o['chord']
            else:
                music21_structure = o['element']
            o['element'].style.color = INTERP[diff]
            f = music21.articulations.Fingering(finger)
            music21_structure.articulations = [f] + music21_structure.articulations
    sf.write('mxl.pdf', fp=path_to_save)

# def replicate_embeddings(clf, x, x_lengths):
#     h, h_last = clf.model.enc(x, x_lengths)
#     # o_attn = clf['attn'](h, h_last)
#     h_last.transpose_(1, 0)
#     # Shape: B x 1 x D_out
#     # Calculate attentional context
#     h.transpose_(1, 2)
#     attention_weights = F.softmax(clf.model.attn.model.attn_ctx(h_last) @ h, dim=0)
#     return attention_weights

def replicate_embeddings(clf, x, x_lengths):
    x_packed = packer(x, x_lengths.cpu(), batch_first=True)

    # Encode
    output, _ = clf.gru1(x_packed)
    output, _ = clf.gru2(output)
    output, hidden = clf.gru3(output)

    # Pass to attention with the original padding
    output_padded, _ = padder(output, batch_first=True)

    e = torch.bmm(clf.attention.w(output_padded), hidden[-1:].permute(1, 2, 0))
    attention_weights = e.softmax(dim=1).detach().cpu().numpy()
    start = 0
    end = 8
    width = end - start

    return np.around((attention_weights - attention_weights.min())/attention_weights.ptp() * width + start).astype(int).squeeze().tolist()

def prediction_torch(matrix, model_path):
    n_features = 88 if "note" in model_path else 10
    n_grades = 3
    # Create a DeepGRU neural network model
    device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
    # load the model
    model = DeepGRU(n_features, n_grades, device=device)
    model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))
    model.eval()
    # convert to torch
    # [b, t, values]
    input = torch.tensor([matrix])
    input_lengths = torch.tensor([input.shape[1]])
    y = model(input, input_lengths)
    y_pred = torch.argmax(y, dim=1)
    print(y_pred)

    attention_weights = replicate_embeddings(model, input, input_lengths)
    return attention_weights

def load_split(split):
    s = load_json("mikrokosmos/splits.json")[str(split)]
    return {"test": s['ids_test'], "train": s['ids_train']}

def load_split_basename(split):
    s = load_json("mikrokosmos/splits.json")[split]
    ans = {"test": [int(os.path.splitext(os.path.basename(x))[0]) for x in s['ids_test']],
            "train": [int(os.path.splitext(os.path.basename(x))[0]) for x in s['ids_train']]}
    ans['train'].sort()
    ans['test'].sort()
    return ans

def get_feature_representation(rep):
    if rep == "note":
        ans = notes_piece
    elif rep == "finger":
        ans = finger_piece
    elif rep == "finger_nakamura":
        ans = finger_nakamura_piece
    elif rep == "velocity":
        ans = velocity_piece
    elif rep == "prob":
        ans = prob_piece
    else:
        raise "bad representation"
    return ans

def export_feedback(split):
    pieces_subsets = load_split(split)
    # pieces_subsets = {'train': ["mikrokosmos/musicxml/69.xml"], 'test': []}
    for appr in ["xgboost"]:  # , "deepgru"
        for subset in ["train", "test"]:
            pieces = pieces_subsets[subset]
            for piece in pieces:
                for rep in ["note", "finger", "finger_nakamura", "prob", "velocity"]:  # ["note", "finger", "finger_nakamura", "velocity", "prob"]:
                    # variables
                    output = f'feedback/{split}/{appr}/{rep}/{subset}'
                    if not os.path.exists(output):
                        os.makedirs(output)
                    if '/1.' in piece or os.path.exists(os.path.join(output, os.path.splitext(os.path.basename(piece))[0] + '.musicxml')):
                        continue
                    if appr == 'deepgru':
                        model = f"results/{appr}/split:{split}_epoch:20_rep_{rep}.pkl"
                    else:
                        model = f"results/{appr}/rep_{rep}/w9/{split}.pkl"

                    # load piece with representation path, grade, path_alias, xml
                    feature_representation = get_feature_representation(rep)
                    matrix, onsets = feature_representation(piece, "mikro2" if rep not in ["finger_nakamura", "prob"] else "nak", piece)

                    # load the model
                    if appr == "deepgru":
                        prediction = prediction_torch(matrix, model)
                    else:
                        windows = salamizer(matrix, 9)
                        clf = xgboost.XGBClassifier()
                        clf.load_model(model)
                        prediction = clf.predict(windows)
                        # else:
                        #     continue

                    # get the values per onset
                    onset_difficulty = get_onset_difficulty(prediction, onsets)

                    # save the output
                    save_score_difficulty("mikro2", output, piece, onset_difficulty, rep, appr)

def update_json():
    structure = {}
    for d in os.listdir('./feedback'):
        if os.path.isdir('./feedback/' + d):
            structure[d] = load_split_basename(d)

    save_json(structure, "feedback_structure.json")

def save_midis():
    if not os.path.exists('mikrokosmos_midis'):
        os.mkdir('mikrokosmos_midis')

    for path, _ in load_json("mikrokosmos/henle_mikrokosmos.json").items():
            path_xml = f"mikrokosmos/musicxml/{path}.xml"
            path_midi = f"mikrokosmos_midis/{path}.mid"

            sc = music21.converter.parse(path_xml)
            sc.write('midi', fp=path_midi)

if __name__ == '__main__':
    assert len(sys.argv) == 2, "Usage: python3 export_feedback [split]"
    export_feedback(33)
    update_json()
    # save_midis()


import tempfile
from pathlib import Path
import shutil
import os
from concurrent.futures import ThreadPoolExecutor

import pytest
import numpy as np
from PyQt5 import QtTest, QtCore, QtGui
import pyperclip
from vtkmodules.vtkRenderingCore import vtkCoordinate
from motmot import geometry
from pangolin import Palmer
import vtkplotlib as vpl
import tomial_tooth_collection_api

from tomial_clicky_tooth._qapp import app
from tomial_clicky_tooth import _csv_io
from tomial_clicky_tooth._ui import UI
from tests import xvfb_size, select_file, CloseBlockingDialog, \
    ChooseMessageBoxButton
from tests.test_csv import INVALID_CSVs, assert_text_equivalent

pytestmark = pytest.mark.order(5)

def screen_coordinates(figure, world):
    """Convert 3D coordinates to 2D pixel numbers in the model renderer."""
    c = vtkCoordinate()
    c.SetCoordinateSystemToWorld()
    c.SetValue(*world)
    x, y = map(int, c.GetComputedDoubleViewportValue(figure.renderer))
    _x, _y = figure.render_size
    return x, _y - y

def click(self, point, button="L"):
    """Click on a 3D point in the renderer."""
    click_2d(self, *screen_coordinates(self.clicker, point), button=button)

def click_2d(self, x, y, button="L"):
    """Click on a given pixel in the model renderer."""
    button = {
        "L": QtCore.Qt.MouseButton.LeftButton,
        "R": QtCore.Qt.MouseButton.RightButton
    }[button]
    QtTest.QTest.mouseClick(self.clicker.vtkWidget, button,
                            QtCore.Qt.KeyboardModifier.NoModifier,
                            QtCore.QPoint(x, y))

def highlight(mesh_plot, point):
    """Color the area around a given point on the mesh.

    This is purely to make it easier to see where some other event (such as a
    mouse click) should be happening. It has no functional purpose.

    """
    mesh_plot.scalars = -geometry.magnitude_sqr(mesh_plot.vertices -
                                                point).clip(max=16)
    mesh_plot.cmap = ["w", "r"]

def tips(mesh, occlusal):
    """Find all locally high points on a mesh.

    This are ideal for clicking on because they are guaranteed to be visible
    when the model is viewed from above.

    """
    heights = occlusal(mesh.vertices)
    ids = mesh.local_maxima(heights, boundaries=False)
    ids = ids[np.argsort(heights[ids])[::-1]]
    return mesh.vertices[ids]

def remove_close(points):
    """Pick out 8 points which are at least 10mm apart from each other."""
    out = np.empty((8, 3), dtype=points.dtype)
    points = iter(points)
    for i in range(8):
        for point in points:
            if (geometry.magnitude_sqr(out[:i] - point) > 100).all():
                out[i] = point
                break
        else:
            raise
    return out

def select_rows(self, *row_numbers):
    """Set the row selection for the UI's table."""
    model = self.table.table.selectionModel()
    selection = QtCore.QItemSelection()
    index = model.model().index
    for i in row_numbers:
        selection.select(index(i, 0), index(i, self.table.shape[1] - 1))
    model.select(selection, model.SelectCurrent)

def action(menu, name):
    """Find a button in a menu of the menu bar."""
    available = [i.text().replace("&", "") for i in menu.actions()]
    for (i, text) in enumerate(available):
        if name.replace("&", "") == text:
            return menu.actions()[i]
    raise ValueError(f"Action named '{name}' not found in {available}")

def test_clicking():
    """Test left and right clicking on the renderer."""
    self = UI(["a", "b", "c", "d"])
    self.show()

    self._open_model(tomial_tooth_collection_api.model("1L"))
    app.processEvents()
    assert np.array(self.table).shape == (4, 3)
    assert np.isnan(np.array(self.table)).all()

    points = remove_close(
        tips(self.clicker.mesh, self.clicker.odometry.occlusal))
    highlight(self.clicker.mesh_plot, points[0])

    assert self.table.highlighted_rows() == [0]
    click(self, points[0])
    assert self.table.highlighted_rows() == [1]
    assert np.allclose(self.table[0], points[0], atol=1)
    assert np.allclose(self.clicker.markers[0].point, points[0], atol=1)

    highlight(self.clicker.mesh_plot, points[1])
    click(self, points[1])
    assert self.table.highlighted_rows() == [2]
    assert np.allclose(self.table[1], points[1], atol=1)
    assert np.allclose(self.clicker.markers[1].point, points[1], atol=1)

    highlight(self.clicker.mesh_plot, points[3])
    select_rows(self, 3)
    click(self, points[3])
    assert self.table.highlighted_rows() == [3]
    assert np.allclose(self.table[3], points[3], atol=1)
    assert np.allclose(self.clicker.markers[3].point, points[3], atol=1)

    highlight(self.clicker.mesh_plot, points[4])
    click(self, points[4])
    assert self.table.highlighted_rows() == [3]
    assert np.allclose(self.table[3], points[4], atol=1)
    assert np.allclose(self.clicker.markers[3].point, points[4], atol=1)

    select_rows(self)
    highlight(self.clicker.mesh_plot, points[5])
    click(self, points[5])
    assert self.table.highlighted_rows() == [2]
    assert np.allclose(self.table[2], points[5], atol=1)
    assert np.allclose(self.clicker.markers[2].point, points[5], atol=1)

    self.clicker.mesh_plot.scalars = None

    click(self, points[5], button="R")
    assert sorted(self.clicker.markers.keys()) == [0, 1, 3]
    click(self, points[4], button="R")
    assert sorted(self.clicker.markers.keys()) == [0, 1]
    click(self, points[0], button="R")
    assert sorted(self.clicker.markers.keys()) == [1]
    click(self, points[0], button="R")
    assert sorted(self.clicker.markers.keys()) == [1]

    # Clicking on empty space should do nothing.
    click_2d(self, 10, 10)
    assert sorted(self.clicker.markers.keys()) == [1]
    click_2d(self, 10, 10, button="R")
    assert sorted(self.clicker.markers.keys()) == [1]

    self.close()

def test_highlight():
    """Test highlighting table rows highlights the corresponding placed markers
    in the renderer."""
    self = UI(Palmer.range("LL5", "LR5"),
              tomial_tooth_collection_api.model("1L"))
    self.show()
    points = remove_close(
        tips(self.clicker.mesh, self.clicker.odometry.occlusal))
    self.points = points

    assert self.clicker.markers[0].color == (.9, 0, 0)
    assert self.clicker.markers[1].color == (0, 0, 0)
    select_rows(self, 0)
    assert self.clicker.markers[0].color == (.9, 0, 0)
    assert self.clicker.markers[1].color == (0, 0, 0)
    select_rows(self, 1)
    assert self.clicker.markers[0].color == (0, 0, 0)
    assert self.clicker.markers[1].color == (.9, 0, 0)
    select_rows(self, 2)
    assert self.clicker.markers[0].color == (0, 0, 0)
    assert self.clicker.markers[1].color == (0, 0, 0)

    select_rows(self, 0, 2, 1)
    assert self.clicker.markers[0].color == (.9, 0, 0)
    assert self.clicker.markers[1].color == (.9, 0, 0)

    self.close()

def test_delete():
    """Test the delete and clear all buttons."""
    self = UI(Palmer.range("LL4", "LR4"),
              tomial_tooth_collection_api.model("1L"))
    self.show()
    points = remove_close(
        tips(self.clicker.mesh, self.clicker.odometry.occlusal))
    self.points = points

    select_rows(self, 0, 1, 2)
    self.table.delete_button.click()
    for i in range(3):
        assert i not in self.clicker.markers
    for i in range(3, 8):
        assert self.clicker.markers[i]
    assert self.table[:3] == [None] * 3
    assert (self.table[3:] == points[3:]).all()

    self.table.clear_all_button.click()
    assert self.table[:] == [None] * 8
    assert self.clicker.markers == {}
    assert len(self.clicker.plots) == 1

    self.close()

@pytest.mark.parametrize("long_names", [False, True])
@pytest.mark.parametrize("show", ["show", "showMaximized", "showFullScreen"])
def test_table_layout(long_names, show):
    """Ensure that the table reactively takes as much horizontal space as it
    needs to avoid needing a horizontal scroll bar, but no more.

    """
    names = Palmer.range()
    if long_names:
        names = [
            f"The bit on the {i} where the landmark needs to go" for i in names
        ]

    self = UI(names)
    self._open_model(tomial_tooth_collection_api.model("1L"))

    getattr(self, show)()
    # Xvfb has no concept of maximised/fullscreen mode and Qt can't detect its
    # screen size so manually set it to use the full screen.
    if xvfb_size():
        if show == "showMaximized":
            return
        if show == "showFullScreen":
            self.resize(*xvfb_size())

    # If not using maximised/fullscreen mode, make the window tiny.
    if show == "show":
        self.resize(600, 400)
    target_width = self.table.table.sizeHint().width()

    for i in range(3):
        app.processEvents()
        if self.table.table.width() == target_width:
            break
    assert self.table.table.width() == target_width

    assert not self.table.table.horizontalScrollBar().isVisible()
    assert target_width < self.table.width() < target_width + 50

    # The table should expand horizontally a bit when landmarks are added to fit
    # the coordinates.
    old_size = self.table.size()
    for point in self.clicker.mesh.vectors[[1000, 20000], 0]:
        self.clicker.spawn_marker(point)
    app.processEvents()
    assert self.table.width() > old_size.width()
    assert self.table.height() == old_size.height()
    assert not self.table.table.horizontalScrollBar().isVisible()

    self.close()

def test_paste():
    """Test pasting points into the landmark table."""
    self = UI(Palmer.range(0, "LR5"))
    self.show()
    app.processEvents()

    # Basic paste one point.
    pyperclip.copy("123\t456\t789")
    action(self.menu_bar["&Edit"], "Paste").trigger()
    assert self.table[0] == (123, 456, 789)
    assert self.clicker.markers[0].point == (123, 456, 789)

    # Populate the table.
    points = np.c_[np.arange(5) * 4, np.random.random(5), np.random.random(5)]
    self.points = points
    self.clicker.reset_camera()
    self.clicker.update()
    assert np.all(self.points == points)

    # Test that various invalid CSVs don't write to the table before realising
    # that they are invalid.
    for text in INVALID_CSVs:
        assert _csv_io.parse_points(text) is None
        pyperclip.copy(text)
        self.table.paste()
        assert np.all(self.points == points)

    # Ensure that trailing newlines don't cause the next point to be overwritten
    # with a blank.
    self.table.table.selectRow(3)
    pyperclip.copy("0,0,0\n")
    self.table.paste()
    assert self.table[3] == (0, 0, 0)
    assert np.all(self.clicker.markers[4].point == points[4])

    # Paste overwriting all values, including replacing points with blanks.
    pyperclip.copy(",,\r\n1, .5, .3\r\n,,\r\n3,.2,.6\r\n,,")
    self.table.table.selectRow(0)
    self.table.paste()
    assert len(self.clicker.markers) == 2
    assert self.clicker.markers[1].point == (1, .5, .3)
    assert self.clicker.markers[3].point == (3, .2, .6)

    # Paste starting from the middle of the table so that the last two values in
    # the clipboard don't fit into the table and should be discarded.
    self.table.table.selectRow(2)
    self.table.paste()
    assert len(self.clicker.markers) == 2
    assert self.clicker.markers[1].point == (1, .5, .3)
    assert self.clicker.markers[3].point == (1, .5, .3)

    # With nothing selected, paste starting from the top.
    select_rows(self)
    self.table.paste()
    assert len(self.clicker.markers) == 2
    assert self.clicker.markers[1].point == (1, .5, .3)
    assert self.clicker.markers[3].point == (3, .2, .6)

    self.close()

def test_copy():
    """Test the copy and cut buttons."""
    pyperclip.copy("")

    points = np.c_[np.arange(5) * 5, np.zeros(5), np.ones(5)]
    self = UI(Palmer.range(0, "LR5"), points=points)
    self.show()
    app.processEvents()

    action(self.menu_bar["&Edit"], "Copy").trigger()
    assert pyperclip.paste() == "0.0\t0.0\t1.0"

    select_rows(self, 1, 2, 4)
    self.table.copy()
    assert pyperclip.paste() == "5.0\t0.0\t1.0\n10.0\t0.0\t1.0\n20.0\t0.0\t1.0"

    select_rows(self, 0, 2)
    action(self.menu_bar["&Edit"], "Cut").trigger()
    assert pyperclip.paste() == "0.0\t0.0\t1.0\n10.0\t0.0\t1.0"
    assert sorted(self.clicker.markers) == [1, 3, 4]

    self.close()

def test_save():
    points = np.arange(6).reshape((2, 3))
    self = UI(["foo", "bar"], points=points)
    self.show()
    app.processEvents()

    def _test_write(name):
        csv_path = os.path.join(root, name)
        with select_file(csv_path):
            self.table.save_as()
        with open(csv_path) as f:
            assert_text_equivalent(f.read(),
                                   "Landmarks,X,Y,Z\nfoo,0,1,2\nbar,3,4,5\n")

    with tempfile.TemporaryDirectory() as root:
        assert self.table.default_csv_path() == ""
        _test_write("foo.csv")
        with select_file(tomial_tooth_collection_api.model("1L")):
            self.open_model()
        self.points = points
        assert self.table.default_csv_path().name == "1L.csv"
        _test_write("1L.csv")

    self.close()

def key_press(widget, key, modifier=QtCore.Qt.NoModifier):
    """Send keyboard events."""
    event = QtGui.QKeyEvent(QtCore.QEvent.KeyPress, key, modifier)
    app.sendEvent(widget, event)

def test_model_switching():
    """Test iterating through a directory of models."""
    from PyQt5.QtCore import Qt

    # Put several models and a points file in a single directory.
    with tempfile.TemporaryDirectory() as root:
        root = Path(root)
        files = [
            Path(shutil.copy(tomial_tooth_collection_api.model(name), root))
            for name in ["1L", "1U", "2L", "2U"]
        ]

        self = UI(Palmer.range(), path=files[0])
        self.show()
        app.processEvents()

        assert self.clicker.path == files[0]
        self.buttons[1].click()
        assert self.clicker.path != files[0]
        assert self.clicker.path in files
        self.buttons[0].click()
        assert self.clicker.path == files[0]

        key_press(self, Qt.Key_Right)
        assert self.clicker.path != files[0]
        key_press(self, Qt.Key_Left)
        assert self.clicker.path == files[0]

        key_press(self, Qt.Key_Left, Qt.ShiftModifier)
        assert self.clicker.path == files[0]
        key_press(self.table.table, Qt.Key_Up)
        assert self.clicker.path == files[0]
        key_press(self.table.table, Qt.Key_Left, Qt.ShiftModifier)
        assert self.clicker.path == files[0]

        # It shouldn't matter which portion of the UI currently has focus.
        key_press(self.clicker, Qt.Key_Right)
        assert self.clicker.path != files[0]
        key_press(self.clicker.vtkWidget, Qt.Key_Left)
        assert self.clicker.path == files[0]
        key_press(self.table, Qt.Key_Right)
        assert self.clicker.path != files[0]
        key_press(self.table.table, Qt.Key_Left)
        assert self.clicker.path == files[0]

        # Check the orientation is adjusting to each model.
        up = self.clicker.odometry.up
        forwards = self.clicker.odometry.forwards
        while not self.clicker.path.name.startswith("1U"):
            self.buttons[0].click()
        view = vpl.view(fig=self.clicker)
        camera_direction = geometry.UnitVector(
            np.array(view["focal_point"]) - view["camera_position"])
        assert camera_direction(up) > .9
        assert forwards(view["up_view"]) > .9

        # With a landmarks file present, it should be picked up.
        shutil.copy(Path(__file__).with_name("1L.csv"), root)
        while not self.clicker.path.name.startswith("1L"):
            self.switch_model(">")
        assert len(self.clicker.markers) == 14
        self.switch_model(">")
        assert len(self.clicker.markers) == 0

        # Without appropriate thread control, the UI will crash if the user
        # holds either left or right keys for around >30 seconds.
        with ThreadPoolExecutor() as pool:
            clicks = [pool.submit(self.buttons[0].click) for i in range(200)]
            [i.result() for i in clicks]
        app.processEvents()

    # Verify that nothing happens if the STL is deleted.
    assert not self.clicker.path.exists()
    old_points = np.arange(12).reshape((4, 3))
    self.points = old_points
    old_stl_plot = self.clicker.mesh_plot
    self.buttons[0].click()
    assert self.clicker.mesh_plot is old_stl_plot
    assert np.array_equal(self.points[:4], old_points)

    self.close()

def test_custom_files_index(tmpdir):
    """Test a custom implementation of UI.files_index() which iterates over a
    predefined set of files instead of over a directory."""
    files = [
        Path(shutil.copy(tomial_tooth_collection_api.model("1L"), tmpdir)),
        tomial_tooth_collection_api.model("2L"),
    ]

    class CustomFileList(UI):
        def files_index(self):
            try:
                return files, files.index(self.path)
            except ValueError:
                return None, None

        def csv_path(self):
            return tmpdir / f"foo-{self.path.name}.csv"

    self = CustomFileList(Palmer.range(), files[0])
    self.show()

    assert self.path == files[0]
    assert self.model_name_indicator.text() == "1L"
    assert self.model_number_indicator.text() == "(1/2)"
    assert self.table.default_csv_path() == tmpdir / "foo-1L.stl.gz.csv"
    self.table.save()
    assert self.table.default_csv_path().exists()

    self.switch_model(">")
    assert self.path == files[1]
    assert self.model_name_indicator.text() == "2L"
    assert self.model_number_indicator.text() == "(2/2)"

    self.switch_model(">")
    assert self.path == files[0]
    assert self.model_name_indicator.text() == "1L"
    assert self.model_number_indicator.text() == "(1/2)"

    self._open_model(tomial_tooth_collection_api.model("2U"))
    assert self.path == tomial_tooth_collection_api.model("2U")
    assert self.model_name_indicator.text() == "2U"
    assert self.model_number_indicator.text() == ""

def test_no_model_open():
    """Make sure that we can't cause a crash by pressing buttons when no model
    is loaded."""
    self = UI(Palmer.range())
    self.show()
    app.processEvents()

    with select_file(None):
        self.open_model()
    assert self.clicker.path is None

    self.buttons[0].click()

    with select_file(None):
        self.table.save_as()

    self.points = [[1, 2, 3]]
    read, write = os.pipe()
    self.table._save(write)
    with open(read, "rb") as f:
        assert f.readlines()[1].startswith(b"*L8,1.0,2.0,3.0")

    self.close()

def test_invalid_model(tmpdir):
    """Test the various points where an invalid model may be loaded."""

    model, invalid, landmarks = [
        Path(shutil.copy(
            i, tmpdir)) for i in (tomial_tooth_collection_api.model("1L"),
                                  Path(__file__).with_name("invalid.stl"),
                                  Path(__file__).with_name("1L.csv"))
    ]

    # Open an invalid model on startup.
    with CloseBlockingDialog():
        self = UI(Palmer.range(), invalid)
    self.show()
    app.processEvents()
    assert self.clicker.path == invalid
    assert self.clicker.mesh is None
    assert self.clicker.mesh_plot is None

    # Via the open model dialog after opening a functional model.
    self._open_model(model)
    assert self.clicker.path == model
    assert np.isfinite(self.clicker.landmarks[1]).sum() == 3 * 14
    with CloseBlockingDialog():
        self._open_model(invalid)
    assert self.clicker.path == invalid
    assert self.clicker.mesh is None
    assert len(self.clicker.markers) == 0

    # By switching models via the buttons.
    self.buttons[0].click()
    assert self.clicker.path == model
    with CloseBlockingDialog():
        self.buttons[0].click()

    self.close()

def test_show_licenses():
    self = UI(Palmer.range())
    self.show()
    app.processEvents()

    self.show_licenses()
    assert self._licenses.isVisible()

def test_history(tmpdir):
    """Test undo/redo."""
    self = UI(Palmer.range())
    self.show()
    app.processEvents()

    assert len(self._history) == 1
    assert self._history.position == 0
    assert np.isnan(self._history[0]).all()
    assert not self.isWindowModified()

    files = [
        Path(shutil.copy(tomial_tooth_collection_api.model(name), tmpdir))
        for name in ["1L", "1U"]
    ]

    old_history = self._history
    self._open_model(files[0])
    assert self._history is not old_history
    assert len(self._history) == 1
    assert self._history.position == 0
    assert np.isnan(self._history[0]).all()
    assert not self.isWindowModified()

    pyperclip.copy("1\t2\t3\n4\t5\t6")
    action(self.menu_bar["&Edit"], "Paste").trigger()
    assert len(self._history) == 2
    assert self._history.position == 1
    assert np.array_equal(self.points[:2], self._history[1][:2])
    assert self.isWindowModified()

    self.undo()
    assert len(self._history) == 2
    assert np.isnan(self.points).all()
    assert not self.isWindowModified()
    self.undo()
    assert len(self._history) == 2
    assert np.isnan(self.points).all()
    assert not self.isWindowModified()

    self.clicker.spawn_marker((7, 8, 9))
    self.clicker.spawn_marker((10, 11, 12))
    assert len(self._history) == 3
    assert self.isWindowModified()
    self.table._save(tmpdir / "foo.csv")
    assert not self.isWindowModified()

    self.undo()
    self.clicker.spawn_marker((10, 11, 12))
    assert len(self._history) == 3
    assert self.isWindowModified()

    self.undo()
    assert self.isWindowModified()
    self.redo()
    assert sorted(self.clicker.markers) == [0, 2]
    self.redo()
    assert sorted(self.clicker.markers) == [0, 2]
    assert self.isWindowModified()

    self.clicker.spawn_marker((13, 14, 15))
    self.clicker.spawn_marker((16, 17, 18))
    for i in range(10):
        self.undo()
    for i in range(10):
        self.redo()

def test_save_prompt(tmp_path):
    """Test that the *would you like to save before moving on* prompt appears
    when it's supposed to."""
    files = [
        Path(shutil.copy(tomial_tooth_collection_api.model(name), tmp_path))
        for name in ["1L", "1U"]
    ]

    self = UI(Palmer.range(), files[0])
    self.show()
    app.processEvents()

    self.clicker.spawn_marker((7, 8, 9))
    assert self._history.modified
    with ChooseMessageBoxButton("Cancel"):
        self.switch_model(">")
    assert self._history.modified
    assert self.path == files[0]

    with ChooseMessageBoxButton("Save"):
        self.switch_model(">")
    assert not self._history.modified
    assert self.path == files[1]
    assert (tmp_path / "1L.csv").exists()

    self.clicker.spawn_marker((7, 8, 9))
    with ChooseMessageBoxButton("Don't Save"):
        self.switch_model(">")
    assert not self._history.modified
    assert self.path == files[0]
    assert not (tmp_path / "1U.csv").exists()

    self.clicker.spawn_marker((7, 8, 9))
    with ChooseMessageBoxButton("Cancel"):
        self.open_model()
    assert self.path == files[0]
    assert self._history.modified

    with select_file(tmp_path / "foo.csv"):
        with ChooseMessageBoxButton("Save As"):
            self.switch_model(">")
        assert (tmp_path / "foo.csv").exists()
        assert self.path == files[1]


from keras.layers import Dense, Dropout, Activation, BatchNormalization, Input
from keras.models import Model
from keras.optimizers import SGD
from keras.models import load_model
from keras import backend as K
from sklearn.utils import class_weight
from sklearn.metrics import roc_auc_score
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import train_test_split
from sklearn.metrics import log_loss
from get_data import GetData

import json
import pandas as pd
import numpy as np
import torch
import sys
import copy
import tensorflow as tf

from scipy.optimize import minimize

sys.path.append('/content/gdrive/My Drive/Optimizasyon-2')
from jtnn import *
sys.path.append('/content/gdrive/My Drive/Optimizasyon-2')

vocab = [x.strip("\r\n ") for x in open("/content/gdrive/My Drive/Optimizasyon-2/Data/unique_canonical_train_vocab.txt")]
vocab = Vocab(vocab)

hidden_size = 450
latent_size = 56
depth = 3
stereo = True

model_jtvae = JTNNVAE(vocab, hidden_size, latent_size, depth, stereo=stereo)
model_jtvae.load_state_dict(torch.load("/content/gdrive/My Drive/Optimizasyon-2/Model/model.iter-9-6000", map_location=torch.device('cpu')))  # opts.model_path

with open('./L1000CDS_subset.json', 'r') as f:
    L = json.load(f)

obj = GetData(L=L, cell_line='PC3', descriptor='jtvae', n_fold=5, random_state=42, random_genes=False, csv_file='Data/JTVAE_Representations.csv')
x, y, folds = obj.get_up_genes()
trn_x = x.drop(['SMILES'], axis=1).values.astype('float')

scaler = StandardScaler()
scaler.fit(trn_x)

cell_line = 'PC3'
gene_target_up = pd.read_csv('Data/harmonizome_dn_binarized_use_for_up_model.csv')
gene_target_dn = pd.read_csv('Data/harmonizome_up_binarized_use_for_dn_model.csv')

gene_target_up = gene_target_up[gene_target_up['disease1'] == 'Cancer of prostate_Prostate_GSE1413']
gene_target_dn = gene_target_dn[gene_target_dn['disease1'] == 'Cancer of prostate_Prostate_GSE1413']

dis_df = pd.read_csv('jtvae-hastalik-drug/prostate adenocarcinoma.csv') # PC3
model_up = load_model('Model/' + cell_line + '_multi_task_model_up.h5')
model_dn = load_model('Model/' + cell_line + '_multi_task_model_dn.h5')

file_name_up = 'Model/' + cell_line + '_multi_task_gene_list_up.txt'
f = open(file_name_up, 'r')
lines = f.readlines()
gene_list_up = [line.strip() for line in lines]
y_true_up = gene_target_up[gene_list_up].values[0]

file_name_dn = 'Model/' + cell_line + '_multi_task_gene_list_dn.txt'
f = open(file_name_dn, 'r')
lines = f.readlines()
gene_list_dn = [line.strip() for line in lines]
y_true_dn = gene_target_dn[gene_list_dn].values[0]

cikarilan_smi_lst = []
baslangic_smi_lst = []
distance_lst = []
optimized_smi_lst = []

def loss_func(par):
    y_up = np.asarray([x[0][0] for x in model_up.predict(np.asarray(par).reshape(1, -1))])
    y_dn = np.asarray([x[0][0] for x in model_dn.predict(np.asarray(par).reshape(1, -1))])

    loss_up = 1 - y_up[25]
    loss_dn = 1 - y_dn[25]

    return loss_up + loss_dn

def optimize(X_initial):
    X_initial = scaler.transform(X_initial.reshape(1, -1)).reshape(1, -1)

    res = minimize(loss_func, X_initial, method='SLSQP', options={'ftol': 1e-9, 'disp': True})

    return scaler.inverse_transform(res.x)

for idx in range(len(dis_df)):
    arr = np.random.uniform(low=-1.5, high=1.5, size=(56,))
    baslangic_feat = dis_df.values[idx, 3:].astype('float') + arr
    dec_smiles = model_jtvae.reconstruct2(torch.from_numpy(np.asarray([baslangic_feat[0:28]])).float(),
                                          torch.from_numpy(np.asarray([baslangic_feat[28:56]])).float())

    cikarilan_smi = dis_df.values[idx, 2]
    baslangic_smi = dec_smiles
    distance = np.linalg.norm(dis_df.values[idx, 3:].astype('float') - baslangic_feat)
    optimize_edilen_smi = ""
    print(cikarilan_smi, ':', baslangic_smi, ':', distance, ':', optimize_edilen_smi)

    updated_feat = optimize(X_initial=baslangic_feat)
    optimize_edilen_smi = model_jtvae.reconstruct2(torch.from_numpy(np.asarray([updated_feat[0:28]])).float(),
                                                   torch.from_numpy(np.asarray([updated_feat[28:56]])).float())

    cikarilan_smi_lst.append(cikarilan_smi)
    baslangic_smi_lst.append(baslangic_smi)
    distance_lst.append(distance)
    optimized_smi_lst.append(optimize_edilen_smi)

    print(cikarilan_smi, ':', baslangic_smi, ':', distance, ':', optimize_edilen_smi)

df = pd.DataFrame({'cikarilan_smiles': cikarilan_smi_lst,
                   'baslangic_smiles': baslangic_smi_lst,
                   'distance': distance_lst,
                   'optimized_smiles': optimized_smi_lst})

df.to_csv('prostate_pc3_multi_task_scipy_SLSQP_15.csv', index=False)


import argparse
import sys
import os
import cv2
import numpy as np
from utils import *
from plate_recognition_image import *

parser = argparse.ArgumentParser()
parser.add_argument('--model-cfg', type=str, default='./cfg/yolov3-face.cfg',
                    help='path to config file')
parser.add_argument('--model-weights', type=str,
                    default='./model-weights/yolov3-wider_16000.weights',
                    help='path to weights of model'),
parser.add_argument('--output-dir', type=str, default='outputs/',
                    help='path to the output directory')
args = parser.parse_args()

net = cv2.dnn.readNetFromDarknet(args.model_cfg, args.model_weights)
net.setPreferableBackend(cv2.dnn.DNN_BACKEND_OPENCV)
net.setPreferableTarget(cv2.dnn.DNN_TARGET_CPU)

included_extensions = ['jpg','jpeg', 'bmp', 'png', 'gif', 'webp']

file_names = [fn for fn in os.listdir('images/')
              if any(fn.endswith(ext) for ext in included_extensions)]

def _main():

    output_file = ''
    for file_name in file_names:
        filename = 'images/'+file_name
        cap = cv2.VideoCapture(filename)
        output_file = file_name.split('.')[0] + '_yoloface.' + file_name.split('.')[1]

        has_frame, frame = cap.read()
        blob = cv2.dnn.blobFromImage(frame, 1 / 255, (IMG_WIDTH, IMG_HEIGHT),
                                     [0, 0, 0], swapRB=True, crop=False)

        net.setInput(blob)
        outs = net.forward(get_outputs_names(net))
        faces = post_process(frame, outs, CONF_THRESHOLD, NMS_THRESHOLD)
        recognition_plate_number(filename, frame)
        save_state = cv2.imwrite(os.path.join(args.output_dir, output_file), frame.astype(np.uint8))
        if save_state :
            print(f"{output_file} successfuly saved!!!")

    cap.release()
    cv2.destroyAllWindows()

if __name__ == '__main__':
    _main()


# -*- coding:utf-8 -*-
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import sys

from PIL import Image, ImageDraw, ImageFont
from shapely.geometry import Polygon
import cv2
import numpy as np
import pyclipper

class DBPreProcess(object):
    def __init__(self, max_side_len=960):
        self.max_side_len = max_side_len

    def resize_image_type(self, im):
        """
        resize image to a size multiple of 32 which is required by the network
        """
        h, w, _ = im.shape

        resize_w = w
        resize_h = h

        # limit the max side
        if max(resize_h, resize_w) > self.max_side_len:
            if resize_h > resize_w:
                ratio = float(self.max_side_len) / resize_h
            else:
                ratio = float(self.max_side_len) / resize_w
        else:
            ratio = 1.
        resize_h = int(resize_h * ratio)
        resize_w = int(resize_w * ratio)
        if resize_h % 32 == 0:
            resize_h = resize_h
        elif resize_h // 32 <= 1:
            resize_h = 32
        else:
            resize_h = (resize_h // 32 - 1) * 32
        if resize_w % 32 == 0:
            resize_w = resize_w
        elif resize_w // 32 <= 1:
            resize_w = 32
        else:
            resize_w = (resize_w // 32 - 1) * 32
        try:
            if int(resize_w) <= 0 or int(resize_h) <= 0:
                return None, (None, None)
            im = cv2.resize(im, (int(resize_w), int(resize_h)))
        except:
            print(im.shape, resize_w, resize_h)
            sys.exit(0)
        ratio_h = resize_h / float(h)
        ratio_w = resize_w / float(w)
        return im, (ratio_h, ratio_w)

    def normalize(self, im):
        img_mean = [0.485, 0.456, 0.406]
        img_std = [0.229, 0.224, 0.225]
        im = im.astype(np.float32, copy=False)
        im = im / 255
        im -= img_mean
        im /= img_std
        channel_swap = (2, 0, 1)
        im = im.transpose(channel_swap)
        return im

    def __call__(self, im):
        im, (ratio_h, ratio_w) = self.resize_image_type(im)
        im = self.normalize(im)
        im = im[np.newaxis, :]
        return [im, (ratio_h, ratio_w)]

class DBPostProcess(object):
    """
    The post process for Differentiable Binarization (DB).
    """

    def __init__(self, thresh=0.3, box_thresh=0.5, max_candidates=1000):
        self.thresh = thresh
        self.box_thresh = box_thresh
        self.max_candidates = max_candidates
        self.min_size = 3

    def boxes_from_bitmap(self, pred, _bitmap, dest_width, dest_height):
        '''
        _bitmap: single map with shape (1, H, W),
                whose values are binarized as {0, 1}
        '''

        bitmap = _bitmap
        height, width = bitmap.shape

        outs = cv2.findContours((bitmap * 255).astype(np.uint8), cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)
        if len(outs) == 3:
            img, contours, _ = outs[0], outs[1], outs[2]
        elif len(outs) == 2:
            contours, _ = outs[0], outs[1]

        num_contours = min(len(contours), self.max_candidates)
        boxes = np.zeros((num_contours, 4, 2), dtype=np.int16)
        scores = np.zeros((num_contours, ), dtype=np.float32)

        for index in range(num_contours):
            contour = contours[index]
            points, sside = self.get_mini_boxes(contour)
            if sside < self.min_size:
                continue
            points = np.array(points)
            score = self.box_score_fast(pred, points.reshape(-1, 2))
            if self.box_thresh > score:
                continue

            box = self.unclip(points).reshape(-1, 1, 2)
            box, sside = self.get_mini_boxes(box)
            if sside < self.min_size + 2:
                continue
            box = np.array(box)
            if not isinstance(dest_width, int):
                dest_width = dest_width.item()
                dest_height = dest_height.item()

            box[:, 0] = np.clip(np.round(box[:, 0] / width * dest_width), 0, dest_width)
            box[:, 1] = np.clip(np.round(box[:, 1] / height * dest_height), 0, dest_height)
            boxes[index, :, :] = box.astype(np.int16)
            scores[index] = score
        return boxes, scores

    def unclip(self, box, unclip_ratio=2.0):
        poly = Polygon(box)
        distance = poly.area * unclip_ratio / poly.length
        offset = pyclipper.PyclipperOffset()
        offset.AddPath(box, pyclipper.JT_ROUND, pyclipper.ET_CLOSEDPOLYGON)
        expanded = np.array(offset.Execute(distance))
        return expanded

    def get_mini_boxes(self, contour):
        bounding_box = cv2.minAreaRect(contour)
        points = sorted(list(cv2.boxPoints(bounding_box)), key=lambda x: x[0])

        index_1, index_2, index_3, index_4 = 0, 1, 2, 3
        if points[1][1] > points[0][1]:
            index_1 = 0
            index_4 = 1
        else:
            index_1 = 1
            index_4 = 0
        if points[3][1] > points[2][1]:
            index_2 = 2
            index_3 = 3
        else:
            index_2 = 3
            index_3 = 2

        box = [points[index_1], points[index_2], points[index_3], points[index_4]]
        return box, min(bounding_box[1])

    def box_score_fast(self, bitmap, _box):
        h, w = bitmap.shape[:2]
        box = _box.copy()
        xmin = np.clip(np.floor(box[:, 0].min()).astype(np.int), 0, w - 1)
        xmax = np.clip(np.ceil(box[:, 0].max()).astype(np.int), 0, w - 1)
        ymin = np.clip(np.floor(box[:, 1].min()).astype(np.int), 0, h - 1)
        ymax = np.clip(np.ceil(box[:, 1].max()).astype(np.int), 0, h - 1)

        mask = np.zeros((ymax - ymin + 1, xmax - xmin + 1), dtype=np.uint8)
        box[:, 0] = box[:, 0] - xmin
        box[:, 1] = box[:, 1] - ymin
        cv2.fillPoly(mask, box.reshape(1, -1, 2).astype(np.int32), 1)
        return cv2.mean(bitmap[ymin:ymax + 1, xmin:xmax + 1], mask)[0]

    def __call__(self, predictions, ratio_list):
        pred = predictions[:, 0, :, :]
        segmentation = pred > self.thresh

        boxes_batch = []
        for batch_index in range(pred.shape[0]):
            height, width = pred.shape[-2:]
            tmp_boxes, tmp_scores = self.boxes_from_bitmap(pred[batch_index], segmentation[batch_index], width, height)

            boxes = []
            for k in range(len(tmp_boxes)):
                if tmp_scores[k] > self.box_thresh:
                    boxes.append(tmp_boxes[k])
            if len(boxes) > 0:
                boxes = np.array(boxes)

                ratio_h, ratio_w = ratio_list[batch_index]
                boxes[:, :, 0] = boxes[:, :, 0] / ratio_w
                boxes[:, :, 1] = boxes[:, :, 1] / ratio_h

            boxes_batch.append(boxes)
        return boxes_batch

def draw_boxes(image, boxes, scores=None, drop_score=0.5):
    img = image.copy()
    draw = ImageDraw.Draw(img)
    if scores is None:
        scores = [1] * len(boxes)
    for (box, score) in zip(boxes, scores):
        if score < drop_score:
            continue
        draw.line([(box[0][0], box[0][1]), (box[1][0], box[1][1])], fill='red')
        draw.line([(box[1][0], box[1][1]), (box[2][0], box[2][1])], fill='red')
        draw.line([(box[2][0], box[2][1]), (box[3][0], box[3][1])], fill='red')
        draw.line([(box[3][0], box[3][1]), (box[0][0], box[0][1])], fill='red')
        draw.line([(box[0][0] - 1, box[0][1] + 1), (box[1][0] - 1, box[1][1] + 1)], fill='red')
        draw.line([(box[1][0] - 1, box[1][1] + 1), (box[2][0] - 1, box[2][1] + 1)], fill='red')
        draw.line([(box[2][0] - 1, box[2][1] + 1), (box[3][0] - 1, box[3][1] + 1)], fill='red')
        draw.line([(box[3][0] - 1, box[3][1] + 1), (box[0][0] - 1, box[0][1] + 1)], fill='red')
    return img

def get_image_ext(image):
    if image.shape[2] == 4:
        return ".png"
    return ".jpg"


import logging
import numpy as np
import numpy.random as npr

from core.config import cfg
import roi_data.data_utils as data_utils
import utils.blob as blob_utils
import utils.boxes as box_utils

logger = logging.getLogger(__name__)

def get_rpn_blob_names(is_training=True):
    """Blob names used by RPN."""
    # im_info: (height, width, image scale)
    blob_names = ['im_info']
    if is_training:
        # gt boxes: (batch_idx, x1, y1, x2, y2, cls)
        blob_names += ['roidb']
        if cfg.FPN.FPN_ON and cfg.FPN.MULTILEVEL_RPN:
            # Same format as RPN blobs, but one per FPN level
            for lvl in range(cfg.FPN.RPN_MIN_LEVEL, cfg.FPN.RPN_MAX_LEVEL + 1):
                blob_names += [
                    'rpn_labels_int32_wide_fpn' + str(lvl),
                    'rpn_bbox_targets_wide_fpn' + str(lvl),
                    'rpn_bbox_inside_weights_wide_fpn' + str(lvl),
                    'rpn_bbox_outside_weights_wide_fpn' + str(lvl)
                ]
        else:
            # Single level RPN blobs
            blob_names += [
                'rpn_labels_int32_wide',
                'rpn_bbox_targets_wide',
                'rpn_bbox_inside_weights_wide',
                'rpn_bbox_outside_weights_wide'
            ]
    return blob_names

def add_rpn_blobs(blobs, im_scales, roidb):
    """Add blobs needed training RPN-only and end-to-end Faster R-CNN models."""
    if cfg.FPN.FPN_ON and cfg.FPN.MULTILEVEL_RPN:
        # RPN applied to many feature levels, as in the FPN paper
        k_max = cfg.FPN.RPN_MAX_LEVEL
        k_min = cfg.FPN.RPN_MIN_LEVEL
        foas = []
        for lvl in range(k_min, k_max + 1):
            field_stride = 2.**lvl
            anchor_sizes = (cfg.FPN.RPN_ANCHOR_START_SIZE * 2.**(lvl - k_min), )
            anchor_aspect_ratios = cfg.FPN.RPN_ASPECT_RATIOS
            foa = data_utils.get_field_of_anchors(
                field_stride, anchor_sizes, anchor_aspect_ratios
            )
            foas.append(foa)
        all_anchors = np.concatenate([f.field_of_anchors for f in foas])
    else:
        foa = data_utils.get_field_of_anchors(cfg.RPN.STRIDE, cfg.RPN.SIZES,
                                              cfg.RPN.ASPECT_RATIOS)
        all_anchors = foa.field_of_anchors

    for im_i, entry in enumerate(roidb):
        scale = im_scales[im_i]
        im_height = np.round(entry['height'] * scale)
        im_width = np.round(entry['width'] * scale)
        gt_inds = np.where(
            (entry['gt_classes'] > 0) & (entry['is_crowd'] == 0)
        )[0]
        gt_rois = entry['boxes'][gt_inds, :] * scale
        # TODO(rbg): gt_boxes is poorly named;
        # should be something like 'gt_rois_info'
        gt_boxes = blob_utils.zeros((len(gt_inds), 6))
        gt_boxes[:, 0] = im_i  # batch inds
        gt_boxes[:, 1:5] = gt_rois
        gt_boxes[:, 5] = entry['gt_classes'][gt_inds]
        im_info = np.array([[im_height, im_width, scale]], dtype=np.float32)
        blobs['im_info'].append(im_info)

        # Add RPN targets
        if cfg.FPN.FPN_ON and cfg.FPN.MULTILEVEL_RPN:
            # RPN applied to many feature levels, as in the FPN paper
            rpn_blobs = _get_rpn_blobs(
                im_height, im_width, foas, all_anchors, gt_rois
            )
            for i, lvl in enumerate(range(k_min, k_max + 1)):
                for k, v in rpn_blobs[i].items():
                    blobs[k + '_fpn' + str(lvl)].append(v)
        else:
            # Classical RPN, applied to a single feature level
            rpn_blobs = _get_rpn_blobs(
                im_height, im_width, [foa], all_anchors, gt_rois
            )
            for k, v in rpn_blobs.items():
                blobs[k].append(v)

    for k, v in blobs.items():
        if isinstance(v, list) and len(v) > 0:
            blobs[k] = np.concatenate(v)

    valid_keys = [
        'has_visible_keypoints', 'boxes', 'segms', 'seg_areas', 'gt_classes',
        'gt_overlaps', 'is_crowd', 'box_to_gt_ind_map', 'gt_keypoints'
    ]
    minimal_roidb = [{} for _ in range(len(roidb))]
    for i, e in enumerate(roidb):
        for k in valid_keys:
            if k in e:
                minimal_roidb[i][k] = e[k]
    # blobs['roidb'] = blob_utils.serialize(minimal_roidb)
    blobs['roidb'] = minimal_roidb

    # Always return valid=True, since RPN minibatches are valid by design
    return True

def _get_rpn_blobs(im_height, im_width, foas, all_anchors, gt_boxes):
    total_anchors = all_anchors.shape[0]
    straddle_thresh = cfg.TRAIN.RPN_STRADDLE_THRESH

    if straddle_thresh >= 0:
        # Only keep anchors inside the image by a margin of straddle_thresh
        # Set TRAIN.RPN_STRADDLE_THRESH to -1 (or a large value) to keep all
        # anchors
        inds_inside = np.where(
            (all_anchors[:, 0] >= -straddle_thresh) &
            (all_anchors[:, 1] >= -straddle_thresh) &
            (all_anchors[:, 2] < im_width + straddle_thresh) &
            (all_anchors[:, 3] < im_height + straddle_thresh)
        )[0]
        # keep only inside anchors
        anchors = all_anchors[inds_inside, :]
    else:
        inds_inside = np.arange(all_anchors.shape[0])
        anchors = all_anchors
    num_inside = len(inds_inside)

    logger.debug('total_anchors: %d', total_anchors)
    logger.debug('inds_inside: %d', num_inside)
    logger.debug('anchors.shape: %s', str(anchors.shape))

    # Compute anchor labels:
    # label=1 is positive, 0 is negative, -1 is don't care (ignore)
    labels = np.empty((num_inside, ), dtype=np.int32)
    labels.fill(-1)
    if len(gt_boxes) > 0:
        # Compute overlaps between the anchors and the gt boxes overlaps
        anchor_by_gt_overlap = box_utils.bbox_overlaps(anchors, gt_boxes)
        # Map from anchor to gt box that has highest overlap
        anchor_to_gt_argmax = anchor_by_gt_overlap.argmax(axis=1)
        # For each anchor, amount of overlap with most overlapping gt box
        anchor_to_gt_max = anchor_by_gt_overlap[np.arange(num_inside),
                                                anchor_to_gt_argmax]

        # Map from gt box to an anchor that has highest overlap
        gt_to_anchor_argmax = anchor_by_gt_overlap.argmax(axis=0)
        # For each gt box, amount of overlap with most overlapping anchor
        gt_to_anchor_max = anchor_by_gt_overlap[
            gt_to_anchor_argmax,
            np.arange(anchor_by_gt_overlap.shape[1])
        ]
        # Find all anchors that share the max overlap amount
        # (this includes many ties)
        anchors_with_max_overlap = np.where(
            anchor_by_gt_overlap == gt_to_anchor_max
        )[0]

        # Fg label: for each gt use anchors with highest overlap
        # (including ties)
        labels[anchors_with_max_overlap] = 1
        # Fg label: above threshold IOU
        labels[anchor_to_gt_max >= cfg.TRAIN.RPN_POSITIVE_OVERLAP] = 1

    # subsample positive labels if we have too many
    num_fg = int(cfg.TRAIN.RPN_FG_FRACTION * cfg.TRAIN.RPN_BATCH_SIZE_PER_IM)
    fg_inds = np.where(labels == 1)[0]
    if len(fg_inds) > num_fg:
        disable_inds = npr.choice(
            fg_inds, size=(len(fg_inds) - num_fg), replace=False
        )
        labels[disable_inds] = -1
    fg_inds = np.where(labels == 1)[0]

    # subsample negative labels if we have too many
    # (samples with replacement, but since the set of bg inds is large most
    # samples will not have repeats)
    num_bg = cfg.TRAIN.RPN_BATCH_SIZE_PER_IM - np.sum(labels == 1)
    bg_inds = np.where(anchor_to_gt_max < cfg.TRAIN.RPN_NEGATIVE_OVERLAP)[0]
    if len(bg_inds) > num_bg:
        enable_inds = bg_inds[npr.randint(len(bg_inds), size=num_bg)]
        labels[enable_inds] = 0
    bg_inds = np.where(labels == 0)[0]

    bbox_targets = np.zeros((num_inside, 4), dtype=np.float32)
    bbox_targets[fg_inds, :] = data_utils.compute_targets(
        anchors[fg_inds, :], gt_boxes[anchor_to_gt_argmax[fg_inds], :]
    )

    # Bbox regression loss has the form:
    #   loss(x) = weight_outside * L(weight_inside * x)
    # Inside weights allow us to set zero loss on an element-wise basis
    # Bbox regression is only trained on positive examples so we set their
    # weights to 1.0 (or otherwise if config is different) and 0 otherwise
    bbox_inside_weights = np.zeros((num_inside, 4), dtype=np.float32)
    bbox_inside_weights[labels == 1, :] = (1.0, 1.0, 1.0, 1.0)

    # The bbox regression loss only averages by the number of images in the
    # mini-batch, whereas we need to average by the total number of example
    # anchors selected
    # Outside weights are used to scale each element-wise loss so the final
    # average over the mini-batch is correct
    bbox_outside_weights = np.zeros((num_inside, 4), dtype=np.float32)
    # uniform weighting of examples (given non-uniform sampling)
    num_examples = np.sum(labels >= 0)
    bbox_outside_weights[labels == 1, :] = 1.0 / num_examples
    bbox_outside_weights[labels == 0, :] = 1.0 / num_examples

    # Map up to original set of anchors
    labels = data_utils.unmap(labels, total_anchors, inds_inside, fill=-1)
    bbox_targets = data_utils.unmap(
        bbox_targets, total_anchors, inds_inside, fill=0
    )
    bbox_inside_weights = data_utils.unmap(
        bbox_inside_weights, total_anchors, inds_inside, fill=0
    )
    bbox_outside_weights = data_utils.unmap(
        bbox_outside_weights, total_anchors, inds_inside, fill=0
    )

    # Split the generated labels, etc. into labels per each field of anchors
    blobs_out = []
    start_idx = 0
    for foa in foas:
        H = foa.field_size
        W = foa.field_size
        A = foa.num_cell_anchors
        end_idx = start_idx + H * W * A
        _labels = labels[start_idx:end_idx]
        _bbox_targets = bbox_targets[start_idx:end_idx, :]
        _bbox_inside_weights = bbox_inside_weights[start_idx:end_idx, :]
        _bbox_outside_weights = bbox_outside_weights[start_idx:end_idx, :]
        start_idx = end_idx

        # labels output with shape (1, A, height, width)
        _labels = _labels.reshape((1, H, W, A)).transpose(0, 3, 1, 2)
        # bbox_targets output with shape (1, 4 * A, height, width)
        _bbox_targets = _bbox_targets.reshape(
            (1, H, W, A * 4)).transpose(0, 3, 1, 2)
        # bbox_inside_weights output with shape (1, 4 * A, height, width)
        _bbox_inside_weights = _bbox_inside_weights.reshape(
            (1, H, W, A * 4)).transpose(0, 3, 1, 2)
        # bbox_outside_weights output with shape (1, 4 * A, height, width)
        _bbox_outside_weights = _bbox_outside_weights.reshape(
            (1, H, W, A * 4)).transpose(0, 3, 1, 2)
        blobs_out.append(
            dict(
                rpn_labels_int32_wide=_labels,
                rpn_bbox_targets_wide=_bbox_targets,
                rpn_bbox_inside_weights_wide=_bbox_inside_weights,
                rpn_bbox_outside_weights_wide=_bbox_outside_weights
            )
        )
    return blobs_out[0] if len(blobs_out) == 1 else blobs_out


# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

import argparse
import logging
import os

import joblib
import numpy as np

from examples.textless_nlp.gslm.speech2unit.clustering.utils import get_audio_files
from examples.textless_nlp.gslm.speech2unit.pretrained.utils import get_features

def get_logger():
    log_format = "[%(asctime)s] [%(levelname)s]: %(message)s"
    logging.basicConfig(format=log_format, level=logging.INFO)
    logger = logging.getLogger(__name__)
    return logger

def get_parser():
    parser = argparse.ArgumentParser(
        description="Quantize using K-means clustering over acoustic features."
    )
    parser.add_argument(
        "--feature_type",
        type=str,
        choices=["logmel", "hubert", "w2v2", "cpc"],
        default=None,
        required=True,
        help="Acoustic feature type",
    )
    parser.add_argument(
        "--kmeans_model_path",
        type=str,
        required=True,
        help="K-means model file path to use for inference",
    )
    parser.add_argument(
        "--manifest_path",
        type=str,
        default=None,
        help="Manifest file containing the root dir and file names",
    )
    parser.add_argument(
        "--checkpoint_path",
        type=str,
        help="Pretrained model checkpoint",
    )
    parser.add_argument(
        "--layer",
        type=int,
        help="The layer of the pretrained model to extract features from",
        default=-1,
    )
    parser.add_argument(
        "--out_dir_path",
        required=True,
        type=str,
        help="File path of quantized output.",
    )
    parser.add_argument(
        "--extension", type=str, default=".flac", help="Features file path"
    )
    return parser

def one_hot(feat, n_clusters):
    return np.eye(n_clusters)[feat]

def main(args, logger):
    # Feature extraction
    logger.info(f"Extracting {args.feature_type} acoustic features...")
    features_batch = get_features(
        feature_type=args.feature_type,
        checkpoint_path=args.checkpoint_path,
        layer=args.layer,
        manifest_path=args.manifest_path,
        sample_pct=1.0,
        flatten=False,
    )
    logger.info(f"Features extracted for {len(features_batch)} utterances.\n")
    logger.info(f"Dimensionality of representation = {features_batch[0].shape[1]}")

    logger.info(f"Loading K-means model from {args.kmeans_model_path} ...")
    kmeans_model = joblib.load(open(args.kmeans_model_path, "rb"))
    kmeans_model.verbose = False

    _, fnames, _ = get_audio_files(args.manifest_path)

    os.makedirs(args.out_dir_path, exist_ok=True)
    logger.info(f"Writing quantized features to {args.out_dir_path}")
    for i, feats in enumerate(features_batch):
        pred = kmeans_model.predict(feats)
        emb = one_hot(pred, kmeans_model.n_clusters)
        base_fname = os.path.basename(fnames[i]).rstrip(args.extension)
        output_path = os.path.join(args.out_dir_path, f"{base_fname}.npy")
        with open(output_path, "wb") as f:
            np.save(f, emb)

if __name__ == "__main__":
    parser = get_parser()
    args = parser.parse_args()
    logger = get_logger()
    logger.info(args)
    main(args, logger)


from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os
import random
import math
import json
import threading
import numpy as np
import tensorflow as tf
import h5py

import util

class PluralModel(object):
  def __init__(self, config):
    self.config = config
    self.context_embeddings = util.EmbeddingDictionary(config["context_embeddings"])
    self.head_embeddings = util.EmbeddingDictionary(config["head_embeddings"], maybe_cache=self.context_embeddings)
    self.char_embedding_size = config["char_embedding_size"]
    self.char_dict = util.load_char_dict(config["char_vocab_path"])
    if config["lm_path"]:
      self.lm_file = h5py.File(self.config["lm_path"], "r")
    else:
      self.lm_file = None
    self.lm_layers = self.config["lm_layers"]
    self.lm_size = self.config["lm_size"]
    self.eval_data = None # Load eval data lazily.

    input_props = []
    input_props.append((tf.float32, [None, None, self.context_embeddings.size])) # Context embeddings.
    input_props.append((tf.float32, [None, None, self.head_embeddings.size])) # Head embeddings.
    input_props.append((tf.float32, [None, None, self.lm_size, self.lm_layers])) # LM embeddings.
    input_props.append((tf.int32, [None, None, None])) # Character indices.
    input_props.append((tf.int32, [None])) # Text lengths.
    input_props.append((tf.bool, [])) # Is training.
    input_props.append((tf.int32, [None])) # Gold starts.
    input_props.append((tf.int32, [None])) # Gold ends.
    input_props.append((tf.int32, [None,None])) # antecedents
    input_props.append((tf.int32, [None])) #antecedent_len
    input_props.append((tf.int32,[None])) #anaphors
    input_props.append((tf.bool,[None,None])) # gold labels

    self.queue_input_tensors = [tf.placeholder(dtype, shape) for dtype, shape in input_props]
    dtypes, shapes = zip(*input_props)
    queue = tf.PaddingFIFOQueue(capacity=10, dtypes=dtypes, shapes=shapes)
    self.enqueue_op = queue.enqueue(self.queue_input_tensors)
    self.input_tensors = queue.dequeue()

    self.global_step = tf.Variable(0, name="global_step", trainable=False)
    self.reset_global_step = tf.assign(self.global_step, 0)

    self.predictions, self.loss = self.get_predictions_and_loss(*self.input_tensors)
    learning_rate = tf.train.exponential_decay(self.config["learning_rate"], self.global_step,
                                               self.config["decay_frequency"], self.config["decay_rate"], staircase=True)
    trainable_params = tf.trainable_variables()
    gradients = tf.gradients(self.loss, trainable_params)
    gradients, _ = tf.clip_by_global_norm(gradients, self.config["max_gradient_norm"])
    optimizers = {
      "adam" : tf.train.AdamOptimizer,
      "sgd" : tf.train.GradientDescentOptimizer
    }
    optimizer = optimizers[self.config["optimizer"]](learning_rate)
    self.train_op = optimizer.apply_gradients(zip(gradients, trainable_params), global_step=self.global_step)

  def start_enqueue_thread(self, session):
    train_examples = []
    auxiliary_train_examples = []
    use_teacher_annealing = self.config["use_teacher_annealing"]
    total_plural_doc, use_plural_doc = 0,0
    total_2nd_doc,use_2nd_doc = 0,0
    auxiliary_train_type = self.config["auxiliary_train_type"] #[plural, coref, bridging]

    if self.config["train_path"]: #the main train_path can only be the plural
      for line in open(self.config["train_path"]):
        doc = json.loads(line)
        doc['mode'] = 'plural'
        total_plural_doc+=1
        if len(doc['plurals']) > 0:
          use_plural_doc+=1
          train_examples.append(doc)

    if self.config["auxiliary_train_path"]:
      for line in open(self.config["auxiliary_train_path"]):
        doc = json.loads(line)
        doc['mode'] = auxiliary_train_type
        total_2nd_doc += 1
        if (auxiliary_train_type == 'plural' and len(doc['plurals']) > 0) or \
              (auxiliary_train_type == 'bridging' and len(doc['bridgings'])>0) or \
              auxiliary_train_type == 'coref':
          use_2nd_doc += 1
          auxiliary_train_examples.append(doc)

    print('Find %d plural train documents use %d.' %(total_plural_doc,use_plural_doc))
    if self.config["auxiliary_train_path"]:
      print('Find %d %s train documents use %d from auxiliary train path.' %(total_2nd_doc,auxiliary_train_type,use_2nd_doc))

    if not train_examples:
      train_examples = auxiliary_train_examples
      auxiliary_train_examples = []

    def _enqueue_loop():
      max_step = self.config["max_step"]
      curr_step = 0
      train_ind = 0
      train_2nd_ind = 0

      random.shuffle(train_examples)
      if auxiliary_train_examples:
        random.shuffle(auxiliary_train_examples)
      while True:
        main_train_ratio = 1.0 if not auxiliary_train_examples else \
          (0.5 if not use_teacher_annealing else curr_step*1.0/max_step)

        if random.random() <= main_train_ratio:
          example = train_examples[train_ind]
          train_ind+=1
          if train_ind >= len(train_examples):
            random.shuffle(train_examples)
            train_ind = 0
        else:
          example = auxiliary_train_examples[train_2nd_ind]
          train_2nd_ind+=1
          if train_2nd_ind >= len(auxiliary_train_examples):
            random.shuffle(auxiliary_train_examples)
            train_2nd_ind = 0

        tensorized_example = self.tensorize_example(example, is_training=True)
        feed_dict = dict(zip(self.queue_input_tensors, tensorized_example))
        session.run(self.enqueue_op, feed_dict=feed_dict)
        curr_step+=1
    enqueue_thread = threading.Thread(target=_enqueue_loop)
    enqueue_thread.daemon = True
    enqueue_thread.start()

  def restore(self, session,model_file_name = 'model.max.ckpt'):
    # Don't try to restore unused variables from the TF-Hub ELMo module.
    vars_to_restore = [v for v in tf.global_variables() if "module/" not in v.name]
    saver = tf.train.Saver(vars_to_restore)
    checkpoint_path = os.path.join(self.config["log_dir"], model_file_name)
    print("Restoring from {}".format(checkpoint_path))
    session.run(tf.global_variables_initializer())
    saver.restore(session, checkpoint_path)

  def load_lm_embeddings(self, doc_key):
    if self.lm_file is None:
      return np.zeros([0, 0, self.lm_size, self.lm_layers])
    file_key = doc_key.replace("/", ":")
    group = self.lm_file[file_key]
    num_sentences = len(list(group.keys()))
    sentences = [group[str(i)][...] for i in range(num_sentences)]
    lm_emb = np.zeros([num_sentences, max(s.shape[0] for s in sentences), self.lm_size, self.lm_layers])
    for i, s in enumerate(sentences):
      lm_emb[i, :s.shape[0], :, :] = s
    return lm_emb

  def tensorize_example(self, example, is_training):
    clusters = example["clusters"]
    doc_mode = 'plural' if not 'mode' in example else example['mode']
    gold_mentions = [tuple(m) for m in example['mentions']]
    gold_starts = np.array([s for s,_ in gold_mentions])
    gold_ends = np.array([e for _,e in gold_mentions])
    cluster_ids = [-1] * len(gold_mentions)
    for cluster_id, cluster in enumerate(clusters):
      for mid in cluster:
        cluster_ids[mid] = cluster_id

    sentences = example["sentences"]
    num_words = sum(len(s) for s in sentences)

    max_sentence_length = max(len(s) for s in sentences)
    max_word_length = max(max(max(len(w) for w in s) for s in sentences), max(self.config["filter_widths"]))
    text_len = np.array([len(s) for s in sentences])
    tokens = [[""] * max_sentence_length for _ in sentences]
    context_word_emb = np.zeros([len(sentences), max_sentence_length, self.context_embeddings.size])
    head_word_emb = np.zeros([len(sentences), max_sentence_length, self.head_embeddings.size])
    char_index = np.zeros([len(sentences), max_sentence_length, max_word_length])
    for i, sentence in enumerate(sentences):
      for j, word in enumerate(sentence):
        tokens[i][j] = word
        context_word_emb[i, j] = self.context_embeddings[word]
        head_word_emb[i, j] = self.head_embeddings[word]
        char_index[i, j, :len(word)] = [self.char_dict[c] for c in word]
    tokens = np.array(tokens)

    doc_key = example["doc_key"]

    plural_bridging_cluster_pairs = {} #for bridging and plural
    if doc_mode == 'plural':
      for mid, aid in example['plurals']:
        msid = cluster_ids[mid]
        asid = cluster_ids[aid]
        if not msid in plural_bridging_cluster_pairs:
          plural_bridging_cluster_pairs[msid] = set()
          plural_bridging_cluster_pairs[msid].add(asid)
    elif doc_mode == 'bridging':
      for mid, aid, is_inv in example['bridgings']:
        if is_inv:#for element-of-inv the antecedent is equavalent to plural
          msid = cluster_ids[mid]
          asid = cluster_ids[aid]
        else:
          msid = cluster_ids[aid]
          asid = cluster_ids[mid]
        if not msid in plural_bridging_cluster_pairs:
          plural_bridging_cluster_pairs[msid] = set()
          plural_bridging_cluster_pairs[msid].add(asid)

    num_anaphors = sum(len(cl) for cid, cl in enumerate(clusters) if cid in plural_bridging_cluster_pairs)
    negative_example_rate = 1 if is_training else 0
    if doc_mode != 'coref' and is_training and self.config["ntimes_negative_examples"] > 0:
      negative_example_rate = max(0.2, num_anaphors*self.config["ntimes_negative_examples"]/(len(gold_mentions)-num_anaphors))

    lm_emb = self.load_lm_embeddings(doc_key)
    max_training_words = self.config["max_training_words"] if doc_mode=='plural' else 1000  #for coref/bridging 1000 words maximum

    if is_training and num_words > max_training_words:
      context_word_emb, head_word_emb, lm_emb, char_index, text_len, is_training, gold_starts, gold_ends,np_cluster_ids\
        = self.truncate_example(context_word_emb, head_word_emb, lm_emb, char_index, text_len, is_training, gold_starts, gold_ends, np.array(cluster_ids),max_training_words)
      cluster_ids = [id.item() for id in np_cluster_ids]

    num_mention = gold_starts.shape[0]
    anaphors = np.arange(num_mention)
    num_top_antecedents = min(num_mention, self.config['max_top_antecedents'])
    antecedents = np.zeros([num_mention, num_top_antecedents])
    gold_labels = np.zeros([num_mention, num_top_antecedents], dtype=np.bool)
    antecedents_len = np.zeros([num_mention])
    us_mask = np.ones([num_mention],np.bool)

    for i in xrange(num_mention):
      antecedents_len[i] = min(i, num_top_antecedents)
      cid = cluster_ids[i]
      if doc_mode == 'coref':
        for j in xrange(min(i, num_top_antecedents)):
          ant = i - j - 1
          antecedents[i, j] = ant
          if cluster_ids[i] == cluster_ids[ant]:
            gold_labels[i, j] = True
      elif doc_mode == 'plural':
        us_mask[i] = cid in plural_bridging_cluster_pairs or random.random() < negative_example_rate
        if not us_mask[i]:
          continue
        for j in xrange(min(i, num_top_antecedents)):
          ant = i - j - 1
          antecedents[i, j] = ant
          if cid in plural_bridging_cluster_pairs and cluster_ids[ant] in plural_bridging_cluster_pairs[cid]:
            gold_labels[i, j] = True
      elif doc_mode == 'bridging':
        # for bridging we seach antecedent from both direction as
        # we switch the antecedent and anaphors for element-of relations
        # for element-of-inv we do not switch the position
        us_mask[i] = cid in plural_bridging_cluster_pairs or random.random() < negative_example_rate
        if not us_mask[i]:
          continue
        nbefore = min(i, num_top_antecedents // 2)
        mafter = min(num_mention - i - 1, num_top_antecedents // 2)
        antecedents_len[i] = nbefore + mafter
        for j in xrange(nbefore):
          ant = i - j - 1
          antecedents[i, j] = ant
          if cid in plural_bridging_cluster_pairs and cluster_ids[ant] in plural_bridging_cluster_pairs[cid]:
            gold_labels[i, j] = True
        for j in xrange(mafter):
          ant = i + j + 1
          antecedents[i, j + nbefore] = ant
          if cid in plural_bridging_cluster_pairs and cluster_ids[ant] in plural_bridging_cluster_pairs[cid]:
            gold_labels[i, j + nbefore] = True

    anaphors = anaphors[us_mask]
    antecedents = antecedents[us_mask]
    antecedents_len = antecedents_len[us_mask]
    gold_labels = gold_labels[us_mask]

    return (context_word_emb, head_word_emb, lm_emb, char_index, text_len, is_training, gold_starts, gold_ends, antecedents,antecedents_len,anaphors,gold_labels)

  def truncate_example(self, context_word_emb, head_word_emb, lm_emb, char_index, text_len, is_training, gold_starts, gold_ends, cluster_ids,max_training_words):
    num_words = sum(text_len)
    assert num_words > max_training_words
    num_sentences = context_word_emb.shape[0]
    max_training_sentences = num_sentences

    while num_words > max_training_words:
      max_training_sentences -= 1
      sentence_offset = random.randint(0, num_sentences - max_training_sentences)
      word_offset = text_len[:sentence_offset].sum()
      num_words = text_len[sentence_offset:sentence_offset + max_training_sentences].sum()

    context_word_emb = context_word_emb[sentence_offset:sentence_offset + max_training_sentences, :, :]
    head_word_emb = head_word_emb[sentence_offset:sentence_offset + max_training_sentences, :, :]
    lm_emb = lm_emb[sentence_offset:sentence_offset + max_training_sentences, :, :, :]
    char_index = char_index[sentence_offset:sentence_offset + max_training_sentences, :, :]
    text_len = text_len[sentence_offset:sentence_offset + max_training_sentences]

    gold_spans = np.logical_and(gold_starts >= word_offset, gold_ends < word_offset + num_words)
    gold_starts = gold_starts[gold_spans] - word_offset
    gold_ends = gold_ends[gold_spans] - word_offset
    cluster_ids = cluster_ids[gold_spans]

    return context_word_emb, head_word_emb, lm_emb, char_index, text_len, is_training, gold_starts, gold_ends, cluster_ids

  def get_dropout(self, dropout_rate, is_training):
    return 1 - (tf.to_float(is_training) * dropout_rate)

  def get_predictions_and_loss(self, context_word_emb, head_word_emb, lm_emb, char_index, text_len, is_training, gold_starts, gold_ends, antecedents,antecedents_len,anaphors, gold_labels):
    self.dropout = self.get_dropout(self.config["dropout_rate"], is_training)
    self.lexical_dropout = self.get_dropout(self.config["lexical_dropout_rate"], is_training)
    self.lstm_dropout = self.get_dropout(self.config["lstm_dropout_rate"], is_training)

    num_sentences = tf.shape(context_word_emb)[0]
    max_sentence_length = tf.shape(context_word_emb)[1]

    context_emb_list = [context_word_emb]
    head_emb_list = [head_word_emb]

    if self.config["char_embedding_size"] > 0:
      char_emb = tf.gather(tf.get_variable("char_embeddings", [len(self.char_dict), self.config["char_embedding_size"]]), char_index) # [num_sentences, max_sentence_length, max_word_length, emb]
      flattened_char_emb = tf.reshape(char_emb, [num_sentences * max_sentence_length, util.shape(char_emb, 2), util.shape(char_emb, 3)]) # [num_sentences * max_sentence_length, max_word_length, emb]
      flattened_aggregated_char_emb = util.cnn(flattened_char_emb, self.config["filter_widths"], self.config["filter_size"]) # [num_sentences * max_sentence_length, emb]
      aggregated_char_emb = tf.reshape(flattened_aggregated_char_emb, [num_sentences, max_sentence_length, util.shape(flattened_aggregated_char_emb, 1)]) # [num_sentences, max_sentence_length, emb]
      context_emb_list.append(aggregated_char_emb)
      head_emb_list.append(aggregated_char_emb)

    if self.lm_file:
      lm_emb_size = util.shape(lm_emb, 2)
      lm_num_layers = util.shape(lm_emb, 3)
      with tf.variable_scope("lm_aggregation"):
        self.lm_weights = tf.nn.softmax(tf.get_variable("lm_scores", [lm_num_layers], initializer=tf.constant_initializer(0.0)))
        self.lm_scaling = tf.get_variable("lm_scaling", [], initializer=tf.constant_initializer(1.0))
      flattened_lm_emb = tf.reshape(lm_emb, [num_sentences * max_sentence_length * lm_emb_size, lm_num_layers])
      flattened_aggregated_lm_emb = tf.matmul(flattened_lm_emb, tf.expand_dims(self.lm_weights, 1)) # [num_sentences * max_sentence_length * emb, 1]
      aggregated_lm_emb = tf.reshape(flattened_aggregated_lm_emb, [num_sentences, max_sentence_length, lm_emb_size])
      aggregated_lm_emb *= self.lm_scaling
      context_emb_list.append(aggregated_lm_emb)

    context_emb = tf.concat(context_emb_list, 2) # [num_sentences, max_sentence_length, emb]
    head_emb = tf.concat(head_emb_list, 2) # [num_sentences, max_sentence_length, emb]
    context_emb = tf.nn.dropout(context_emb, self.lexical_dropout) # [num_sentences, max_sentence_length, emb]
    head_emb = tf.nn.dropout(head_emb, self.lexical_dropout) # [num_sentences, max_sentence_length, emb]

    text_len_mask = tf.sequence_mask(text_len, maxlen=max_sentence_length) # [num_sentence, max_sentence_length]

    context_outputs = self.lstm_contextualize(context_emb, text_len, text_len_mask) # [num_words, emb]

    flattened_head_emb = self.flatten_emb_by_sentence(head_emb, text_len_mask) # [num_words]

    mention_emb = self.get_span_emb(flattened_head_emb, context_outputs, gold_starts, gold_ends)

    k=util.shape(antecedents,0)
    c = util.shape(antecedents,1)

    anaphor_emb = tf.gather(mention_emb,anaphors) #[k,emb]
    antecedent_emb = tf.gather(mention_emb, antecedents) # [k, c, emb]

    pair_emb = self.get_pair_embeddings(anaphor_emb, antecedents, antecedent_emb) # [k, c,emb]

    with tf.variable_scope("plural_scores"):
      plural_scores = util.ffnn(pair_emb, self.config["ffnn_depth"], self.config["ffnn_size"], 1, self.dropout) # [k, c, 1]
      plural_scores = tf.squeeze(plural_scores, 2) # [k, c]
      plural_scores = plural_scores + tf.log(tf.sequence_mask(antecedents_len,c,tf.float32))

    dummy_scores = tf.zeros([k, 1])
    dummy_labels = tf.logical_not(tf.reduce_any(gold_labels, 1, keepdims=True))  # [k, 1]

    plural_scores_with_dummy = tf.concat([dummy_scores, plural_scores], 1)
    gold_labels_with_dummy = tf.concat([dummy_labels, gold_labels], 1)

    loss = self.softmax_loss(plural_scores_with_dummy,gold_labels_with_dummy)
    loss = tf.reduce_sum(loss)

    return [plural_scores,antecedents_len,anaphors], loss

  def get_span_emb(self, head_emb, context_outputs, span_starts, span_ends):
    span_emb_list = []

    span_start_emb = tf.gather(context_outputs, span_starts) # [k, emb]
    span_emb_list.append(span_start_emb)

    span_end_emb = tf.gather(context_outputs, span_ends) # [k, emb]
    span_emb_list.append(span_end_emb)

    span_width = tf.minimum(1 + span_ends - span_starts,self.config["max_span_width"]) # [k]

    if self.config["use_features"]:
      span_width_index = span_width - 1 # [k]
      span_width_emb = tf.gather(tf.get_variable("span_width_embeddings", [self.config["max_span_width"], self.config["feature_size"]]), span_width_index) # [k, emb]
      span_width_emb = tf.nn.dropout(span_width_emb, self.dropout)
      span_emb_list.append(span_width_emb)

    if self.config["model_heads"]:
      span_indices = tf.expand_dims(tf.range(self.config["max_span_width"]), 0) + tf.expand_dims(span_starts, 1) # [k, max_span_width]
      span_indices = tf.minimum(util.shape(context_outputs, 0) - 1, span_indices) # [k, max_span_width]
      span_text_emb = tf.gather(head_emb, span_indices) # [k, max_span_width, emb]
      with tf.variable_scope("head_scores"):
        self.head_scores = util.projection(context_outputs, 1) # [num_words, 1]
      span_head_scores = tf.gather(self.head_scores, span_indices) # [k, max_span_width, 1]
      span_mask = tf.expand_dims(tf.sequence_mask(span_width, self.config["max_span_width"], dtype=tf.float32), 2) # [k, max_span_width, 1]
      span_head_scores += tf.log(span_mask) # [k, max_span_width, 1]
      span_attention = tf.nn.softmax(span_head_scores, 1) # [k, max_span_width, 1]
      span_head_emb = tf.reduce_sum(span_attention * span_text_emb, 1) # [k, emb]
      span_emb_list.append(span_head_emb)

    span_emb = tf.concat(span_emb_list, 1) # [k, emb]
    return span_emb # [k, emb]

  def softmax_loss(self, antecedent_scores, antecedent_labels):
    gold_scores = antecedent_scores + tf.log(tf.to_float(antecedent_labels)) # [k, max_ant + 1]
    marginalized_gold_scores = tf.reduce_logsumexp(gold_scores, [1]) # [k]
    log_norm = tf.reduce_logsumexp(antecedent_scores, [1]) # [k]
    return log_norm - marginalized_gold_scores # [k]

  def bucket_distance(self, distances):
    """
    Places the given values (designed for distances) into 10 semi-logscale buckets:
    [0, 1, 2, 3, 4, 5-7, 8-15, 16-31, 32-63, 64+].
    """
    logspace_idx = tf.to_int32(tf.floor(tf.log(tf.to_float(distances))/math.log(2))) + 3
    use_identity = tf.to_int32(distances <= 4)
    combined_idx = use_identity * distances + (1 - use_identity) * logspace_idx
    return tf.clip_by_value(combined_idx, 0, 9)

  def get_pair_embeddings(self, mention_emb, antecedents, antecedent_emb):
    k = util.shape(mention_emb, 0)
    c = util.shape(antecedents, 1)

    feature_emb_list = []
    antecedent_offsets = tf.tile(tf.expand_dims(tf.range(c) + 1, 0), [k, 1]) # [k, c]

    if self.config["use_features"]:
      antecedent_distance_buckets = self.bucket_distance(antecedent_offsets) # [k, c]
      antecedent_distance_emb = tf.gather(tf.get_variable("antecedent_distance_emb", [10, self.config["feature_size"]]), antecedent_distance_buckets) # [k, c]
      feature_emb_list.append(antecedent_distance_emb)

    feature_emb = tf.concat(feature_emb_list, 2) # [k, c, emb]
    feature_emb = tf.nn.dropout(feature_emb, self.dropout) # [k, c, emb]

    target_emb = tf.expand_dims(mention_emb, 1) # [k, 1, emb]
    similarity_emb = antecedent_emb * target_emb # [k, c, emb]
    target_emb = tf.tile(target_emb, [1, c, 1]) # [k, c, emb]

    pair_emb = tf.concat([target_emb, antecedent_emb, similarity_emb, feature_emb], 2) # [k, c, emb]

    return pair_emb

  def flatten_emb_by_sentence(self, emb, text_len_mask):
    num_sentences = tf.shape(emb)[0]
    max_sentence_length = tf.shape(emb)[1]

    emb_rank = len(emb.get_shape())
    if emb_rank  == 2:
      flattened_emb = tf.reshape(emb, [num_sentences * max_sentence_length])
    elif emb_rank == 3:
      flattened_emb = tf.reshape(emb, [num_sentences * max_sentence_length, util.shape(emb, 2)])
    else:
      raise ValueError("Unsupported rank: {}".format(emb_rank))
    return tf.boolean_mask(flattened_emb, tf.reshape(text_len_mask, [num_sentences * max_sentence_length]))

  def lstm_contextualize(self, text_emb, text_len, text_len_mask):
    num_sentences = tf.shape(text_emb)[0]

    current_inputs = text_emb # [num_sentences, max_sentence_length, emb]

    for layer in xrange(self.config["contextualization_layers"]):
      with tf.variable_scope("layer_{}".format(layer)):
        with tf.variable_scope("fw_cell"):
          cell_fw = util.CustomLSTMCell(self.config["contextualization_size"], num_sentences, self.lstm_dropout)
        with tf.variable_scope("bw_cell"):
          cell_bw = util.CustomLSTMCell(self.config["contextualization_size"], num_sentences, self.lstm_dropout)
        state_fw = tf.contrib.rnn.LSTMStateTuple(tf.tile(cell_fw.initial_state.c, [num_sentences, 1]), tf.tile(cell_fw.initial_state.h, [num_sentences, 1]))
        state_bw = tf.contrib.rnn.LSTMStateTuple(tf.tile(cell_bw.initial_state.c, [num_sentences, 1]), tf.tile(cell_bw.initial_state.h, [num_sentences, 1]))

        (fw_outputs, bw_outputs), _ = tf.nn.bidirectional_dynamic_rnn(
          cell_fw=cell_fw,
          cell_bw=cell_bw,
          inputs=current_inputs,
          sequence_length=text_len,
          initial_state_fw=state_fw,
          initial_state_bw=state_bw)

        text_outputs = tf.concat([fw_outputs, bw_outputs], 2) # [num_sentences, max_sentence_length, emb]
        text_outputs = tf.nn.dropout(text_outputs, self.lstm_dropout)
        if layer > 0:
          highway_gates = tf.sigmoid(util.projection(text_outputs, util.shape(text_outputs, 2))) # [num_sentences, max_sentence_length, emb]
          text_outputs = highway_gates * text_outputs + (1 - highway_gates) * current_inputs
        current_inputs = text_outputs

    return self.flatten_emb_by_sentence(text_outputs, text_len_mask)

  def get_plural_pairs(self, plural_scores,antecedent_len,anaphors, cluster_ids, plural_mentions):
    plural_pairs = []
    plural_maps = {}
    for i, mid in enumerate(anaphors):
      if not mid in plural_mentions:
        continue
      score_list = []
      ant_cid_list = set()
      for j in xrange(antecedent_len[i]):
        ant = mid - j -1
        score_list.append((plural_scores[i,j].item(), ant))
      score_list = sorted(score_list,reverse=True)
      for s, ant in score_list:
        ant_cid = cluster_ids[ant]
        if len(ant_cid_list) >=5:
          break
        elif ant_cid in ant_cid_list:
          continue
        elif s > 0:
          ant_cid_list.add(ant_cid)
        elif len(ant_cid_list) < 2:
          ant_cid_list.add(ant_cid)
      plural_maps[mid] = ant_cid_list
      for cid in ant_cid_list:
        plural_pairs.append((mid, cid))
    return set(plural_pairs),plural_maps

  def load_eval_data(self):
    if self.eval_data is None:
      def load_line(line):
        example = json.loads(line)
        return self.tensorize_example(example, is_training=False), example
      with open(self.config["eval_path"]) as f:
        self.eval_data = [load_line(l) for l in f.readlines()]
      print("Loaded {} eval examples.".format(len(self.eval_data)))

  def evaluate(self, session):
    self.load_eval_data()

    tp,fn,fp = 0,0,0
    emp,emt = 0,0

    for example_num, (tensorized_example, example) in enumerate(self.eval_data):
      feed_dict = {i:t for i,t in zip(self.input_tensors, tensorized_example)}
      plural_scores,antecedent_len,anaphors = session.run(self.predictions, feed_dict=feed_dict)

      num_of_mention = len(example['mentions'])
      clusters = example['clusters']
      cluster_ids = [-1] * num_of_mention
      for cid, cl in enumerate(clusters):
        for m in cl:
          cluster_ids[m] = cid

      plural_mentions = set(mid for mid,_ in example['plurals'])
      gold_plural_maps = {}
      for mid, ant in example['plurals']:
        if mid not in gold_plural_maps:
          gold_plural_maps[mid] = set()
        gold_plural_maps[mid].add(cluster_ids[ant])

      gold_plurals = set((mid, cluster_ids[ant]) for mid, ant in example['plurals'])
      pred_plurals, pred_plural_maps = self.get_plural_pairs(plural_scores,antecedent_len,anaphors,cluster_ids,plural_mentions)

      tp += len(gold_plurals & pred_plurals)
      fn += len(gold_plurals - pred_plurals)
      fp += len(pred_plurals - gold_plurals)

      emt += len(gold_plural_maps)
      for mid in gold_plural_maps:
        if len(gold_plural_maps[mid]) == len(pred_plural_maps[mid])==len(gold_plural_maps[mid] & pred_plural_maps[mid]):
          emp+=1

      if example_num % 10 == 0:
        print("Evaluated {}/{} examples.".format(example_num + 1, len(self.eval_data)))

    recall = 0.0 if tp == 0 else float(tp)/(tp+fn)
    precision = 0.0 if tp == 0 else float(tp)/(tp+fp)
    f1 = 0.0 if precision == 0.0 else 2.0*recall*precision/(recall+precision)

    print("Plural F1: {:.2f}%".format(f1*100))
    print("Plural recall: {:.2f}%".format(recall*100))
    print("Plural precision: {:.2f}%".format(precision*100))
    print("Plural exact match accrucy: {:.2f}%".format(emp*100.0/emt))

    summary_dict = {}
    summary_dict["Plural F1"] = f1
    summary_dict["Plural recall"] = recall
    summary_dict["Plural precision"] = precision

    return util.make_summary(summary_dict), f1*100