import streamlit as st
import pandas as pd
import numpy as np
from io import StringIO
from sklearn.feature_extraction.text import CountVectorizer
from sklearn.decomposition import LatentDirichletAllocation, PCA
from sklearn.cluster import KMeans
from sklearn.manifold import TSNE
import re
from russpelling import *
try:
    import maru
except:
    pass
from natasha import Doc, Segmenter, MorphVocab, NewsEmbedding, NewsMorphTagger
import sys
import time
from threading import Thread
from scipy.special import gammaln
from scipy.stats import mode
from scipy.sparse import csr_matrix, lil_matrix
from umap import UMAP
import plotly.express as px
import scipy
from hmmlearn.hmm import CategoricalHMM, GMMHMM
from razdel import tokenize,sentenize
from razdel.substring import Substring
import colorsys
from docx import Document
from docx.shared import RGBColor, Pt
from docx.enum.text import WD_COLOR_INDEX, WD_ALIGN_PARAGRAPH
from soyclustering import SphericalKMeans
from bertopic import BERTopic
from sentence_transformers import SentenceTransformer
import plotly.graph_objects as go
from typing import List, Union
from bertopic.backend._utils import select_backend
import matplotlib.pyplot as plt
from contextualized_topic_models.models.ctm import CombinedTM
from contextualized_topic_models.utils.data_preparation import TopicModelDataPreparation, bert_embeddings_from_list
import datetime
import torch
from contextualized_topic_models.datasets.dataset import CTMDataset
import pyLDAvis as vis
from streamlit import components
from Top2VecNew import Top2VecNew
from matplotlib.colors import rgb2hex
from scipy.optimize import linear_sum_assignment
import math
from navec import Navec
from scipy.ndimage import gaussian_filter1d
from sklearn.metrics.pairwise import cosine_distances
from sklearn.mixture import GaussianMixture


st.set_page_config(layout='wide')
token_pattern = '(?u)\\b\\w+(?:-\\w+)*\\b'
token_pattern_filt = '(?u)\\b[^\\W0-9](?:[-\']?[^\\W0-9])+\\b'
#token_pattern_filt = '[^ ][^ ]+'


class ReturnableThread(Thread):
    def __init__(self, target, args=(), kwargs={}):
        super().__init__(target=target, args=args, kwargs=kwargs)
        self.result = None

    def run(self):
        if self._target:
            self.result = self._target(*self._args, **self._kwargs)
            
def tm_data_prep_fit(self, text_for_contextual, text_for_bow, vectorizer=None):
    if text_for_bow is not None:
        assert len(text_for_contextual) == len(text_for_bow)
    if self.contextualized_model is None:
        raise Exception(
            "A contextualized model must be defined"
        )
    # TODO: this count vectorizer removes tokens that have len = 1, might be unexpected for the users
    self.vectorizer = CountVectorizer() if not vectorizer else vectorizer
    train_bow_embeddings = self.vectorizer.fit_transform(text_for_bow)
    # if the user is passing custom embeddings we don't need to create the embeddings using the model
    train_contextualized_embeddings = bert_embeddings_from_list(
                text_for_contextual,
                sbert_model_to_load=self.contextualized_model,
                max_seq_length=self.max_seq_length,
            )
    self.vocab = self.vectorizer.get_feature_names_out()
    self.id2token = {k: v for k, v in zip(range(0, len(self.vocab)), self.vocab)}
    encoded_labels = None
    return CTMDataset(
        X_contextual=train_contextualized_embeddings,
        X_bow=train_bow_embeddings,
        idx2token=self.id2token,
        labels=encoded_labels,
    )


def gmmfit(gmm, X, w=None, max_iter=1000, eps=1e-6):
    if w is None:
        w = np.ones((X.shape[0]))
    p0 = (gmm.score_samples(X) * w).sum() / w.sum()
    for i in range(max_iter):
        comps = w[:,None] * gmm.predict_proba(X)
        gmm.weights_ = comps.sum(axis=0) / comps.sum()
        
        comps /= comps.sum(axis=0)
        means = np.transpose(comps) @ X
        gmm.means_ = np.where(np.isnan(means), gmm.means_, means)
        
        covars = X[:,None,:] - gmm.means_
        covars = (comps[...,None,None] * (covars[:,:,:,None] * covars[:,:,None,:])).sum(axis=0)
        gmm.covariances_ = np.where(np.isnan(covars), gmm.covariances_, covars)
        
        # КРИТИЧЕСКОЕ ИСПРАВЛЕНИЕ: обновляем precisions_cholesky_
        precisions_chol = np.empty_like(gmm.covariances_)
        for k, cov in enumerate(gmm.covariances_):
            stop = False
            eps = 1e-6
            while not stop:
                val, vec = np.linalg.eig(cov)
                gmm.covariances_[k] = vec.real @ np.diag(np.maximum(val.real,eps)) @ vec.real.transpose()
                try:
                    # Считаем нижнетреугольную матрицу разложения Холецкого
                    cov_chol = np.linalg.cholesky(cov)
                    stop = True
                except np.linalg.LinAlgError:
                    eps *= 10                                               
            # Вычисляем матрицу точности (precision matrix)
            precisions_chol[k] = scipy.linalg.solve_triangular(cov_chol, np.eye(cov_chol.shape[0]), lower=True).T
        gmm.precisions_cholesky_ = precisions_chol
        
        p1 = (gmm.score_samples(X) * w).sum() / w.sum()
        if p1-p0 < eps:
            return
        p0 = p1


def visualize_documents(
    topic_model,
    docs: List[str],
    topics: List[int] = None,
    embeddings: np.ndarray = None,
    reduced_embeddings: np.ndarray = None,
    sample: float = None,
    hide_annotations: bool = False,
    hide_document_hover: bool = False,
    custom_labels: Union[bool, str] = False,
    title: str = "<b>Документы и темы</b>",
    width: int = 1200,
    height: int = 750,
):
    topic_per_doc = topic_model.topics_ if topics is None else topics
    
    # Sample the data to optimize for visualization and dimensionality reduction
    if sample is None or sample > 1:
        sample = 1
        
    indices = []
    for topic in set(topic_per_doc):
        s = np.where(np.array(topic_per_doc) == topic)[0]
        size = len(s) if len(s) < 100 else int(len(s) * sample)
        indices.extend(np.random.choice(s, size=size, replace=False))
    indices = np.array(indices)

    df = pd.DataFrame({"topic": np.array(topic_per_doc)[indices]})
    df["doc"] = [docs[index] for index in indices]
    df["topic"] = [topic_per_doc[index] for index in indices]

    # Extract embeddings if not already done
    if sample is None:
        if embeddings is None and reduced_embeddings is None:
            embeddings_to_reduce = topic_model._extract_embeddings(df.doc.to_list(), method="document")
        else:
            embeddings_to_reduce = embeddings
    else:
        if embeddings is not None:
            embeddings_to_reduce = embeddings[indices]
        elif embeddings is None and reduced_embeddings is None:
            embeddings_to_reduce = topic_model._extract_embeddings(df.doc.to_list(), method="document")

    # Reduce input embeddings
    if reduced_embeddings is None:
        try:
            from umap import UMAP
            umap_model = UMAP(n_neighbors=10, n_components=2, min_dist=0.0, metric="cosine").fit(embeddings_to_reduce)
            embeddings_2d = umap_model.embedding_
        except (ImportError, ModuleNotFoundError):
            raise ModuleNotFoundError(
                "UMAP is required if the embeddings are not yet reduced in dimensionality. Please install it using `pip install umap-learn`."
            )
    elif sample is not None and reduced_embeddings is not None:
        embeddings_2d = reduced_embeddings[indices]
    elif sample is None and reduced_embeddings is not None:
        embeddings_2d = reduced_embeddings
        
    unique_topics = set(topic_per_doc)
    if topics is None:
        topics = unique_topics

    # Combine data
    df["x"] = embeddings_2d[:, 0]
    df["y"] = embeddings_2d[:, 1]

    # Prepare text and names
    if isinstance(custom_labels, str):
        names = [[[str(topic), None]] + topic_model.topic_aspects_[custom_labels][topic] for topic in unique_topics]
        names = ["_".join([label[0] for label in labels[:4]]) for labels in names]
        names = [label if len(label) < 30 else label[:27] + "..." for label in names]
    elif topic_model.custom_labels_ is not None and custom_labels:
        names = [topic_model.custom_labels_[topic + topic_model._outliers] for topic in unique_topics]
    else:
        names = [
            #f"{topic+1}_" + "_".join([word for word, value in topic_model.get_topic(topic)][:3])
            f"{topic+1}_" + "_".join([word for word in st.session_state['topicnames'][topic].split(', ')[:3]])
            for topic in unique_topics
        ]

    # Visualize
    fig = go.Figure()

    # Outliers and non-selected topics
    non_selected_topics = set(unique_topics).difference(topics)
    if len(non_selected_topics) == 0:
        non_selected_topics = [-1]

    selection = df.loc[df.topic.isin(non_selected_topics), :]
    selection["text"] = ""
    selection.loc[len(selection), :] = [
        None,
        None,
        selection.x.mean(),
        selection.y.mean(),
        "Другие документы",
    ]
    
    if False:
        fig.add_trace(
            go.Scattergl(
                x=selection.x,
                y=selection.y,
                hovertext=selection.doc if not hide_document_hover else None,
                hoverinfo="text",
                mode="markers",
                name="other",
                showlegend=False,
                marker=dict(color="#CFD8DC", size=5, opacity=0.5),
            )
    )
    
    # Selected topics
    for name, topic in zip(names, unique_topics):
        if topic in topics and topic != -1:
            selection = df.loc[df.topic == topic, :]
            selection["text"] = ""

            if not hide_annotations:
                selection.loc[len(selection), :] = [
                    None,
                    None,
                    selection.x.mean(),
                    selection.y.mean(),
                    name,
                ]

            fig.add_trace(
                go.Scattergl(
                    x=selection.x,
                    y=selection.y,
                    hovertext=selection.doc if not hide_document_hover else None,
                    hoverinfo="text",
                    text=selection.text,
                    mode="markers",
                    name=name,
                    textfont=dict(
                        size=12,
                    ),
                    marker=dict(color="#CFD8DC", size=5, opacity=0.5) if topic == 0 else dict(size=5, opacity=0.5),
                )
            )

    # Add grid in a 'plus' shape
    x_range = (
        df.x.min() - abs((df.x.min()) * 0.15),
        df.x.max() + abs((df.x.max()) * 0.15),
    )
    y_range = (
        df.y.min() - abs((df.y.min()) * 0.15),
        df.y.max() + abs((df.y.max()) * 0.15),
    )
    fig.add_shape(
        type="line",
        x0=sum(x_range) / 2,
        y0=y_range[0],
        x1=sum(x_range) / 2,
        y1=y_range[1],
        line=dict(color="#CFD8DC", width=2),
    )
    fig.add_shape(
        type="line",
        x0=x_range[0],
        y0=sum(y_range) / 2,
        x1=x_range[1],
        y1=sum(y_range) / 2,
        line=dict(color="#9E9E9E", width=2),
    )
    fig.add_annotation(x=x_range[0], y=sum(y_range) / 2, text="D1", showarrow=False, yshift=10)
    fig.add_annotation(y=y_range[1], x=sum(x_range) / 2, text="D2", showarrow=False, xshift=10)
    
    # Stylize layout
    fig.update_layout(
        template="simple_white",
        title={
            "text": f"{title}",
            "x": 0.5,
            "xanchor": "center",
            "yanchor": "top",
            "font": dict(size=22, color="Black"),
        },
        width=width,
        height=height,
    )
    
    fig.update_xaxes(visible=False)
    fig.update_yaxes(visible=False)
    
    return fig


if False:
    rgb = []
    for i in range((n_topics+1) // 2):
        h, s, v = i / ((n_topics+1)//2), 1.0, 1.0
        r, g, b = colorsys.hsv_to_rgb(h, s, v)
        y, _, _ = colorsys.rgb_to_yiq(r, g, b)
        rgb.append([min(1, 0.6/y*r), min(1, 0.6/y*g), min(1, 0.6/y*b)])
    rgb = rgb + [[0.5*r, 0.5*g, 0.5*b] for r,g,b in rgb]
else:
    rgb = [[0,0,1], [0,0,0.5], [0.5,0,0], [0.5,0.5,0], [0.5,0.5,0.5], [0,0.5,0], [1,0,1], [1,0,0], [0,0.5,0.5], [0.5,0,0.5]]

byword = False
markup_cases = ['HMM','SentHMM']


def modernize_text(text_orig):
    text = text_orig
    for j in re.finditer('i', text.lower()):
        k = j.span()[0]
        if set(text[max(0,k-1):k+2].lower()).intersection('абвгдеёжзийклмнопрстуфхцчшщъыьэюя'):
            text = text[:k] + ('и' if text[k] == 'i' else 'И') + text[k+1:]    
    text = re.sub('чьк', 'чк', text)
    text = re.sub('чьт', 'чт', text)
    text = re.sub('чьп', 'чп', text)
    text = re.sub('чьв', 'чв', text)
    
    text = normalize(text)
    
    text = re.sub('кия\\b', 'кие', text)
    text = re.sub('яго\\b', 'его', text)
    text = re.sub('\\b([Хх])ороше\\b', r'\1орошо', text)
    
    return text

def prepare_vocab(norm_cb, lemma_cb, pr_pbar):
    docs = st.session_state['docs']
    for i in range(len(docs)):
        docs[i]['text'] = modernize_text(docs[i]['text_orig']) if norm_cb else docs[i]['text_orig']
        
    if lemma_cb:
        if 'analyzer' not in st.session_state:
            try:
                st.session_state['analyzer'] = maru.get_analyzer(tagger='rnn', lemmatizer='pymorphy')
            except:
                st.session_state['analyzer'] = {'segmenter': Segmenter(), 'morph_vocab': MorphVocab(), 'emb': NewsEmbedding()}
                st.session_state['analyzer']['morph_tagger'] = NewsMorphTagger(st.session_state['analyzer']['emb'])
    for i in range(len(docs)):
        docs[i]['tokens'],docs[i]['input'] = [],[]
        if not byword:
            docs[i]['sents'] = [s for s in sentenize(docs[i]['text'])]
            chunks = [t for t in tokenize(docs[i]['text'])]
        else:
            docs[i]['sents'] = [Substring(t.span()[0], t.span()[1], t.group()) for t in re.finditer(token_pattern_filt, docs[i]['text'])]
            chunks = docs[i]['sents']
        if lemma_cb:
            try:
                analyzed = st.session_state['analyzer'].analyze([t.text for t in chunks])
                lemmas = [morph.lemma for morph in analyzed]
            except:
                doc = Doc(docs[i]['text'])
                doc.segment(st.session_state['analyzer']['segmenter'])
                doc.tag_morph(st.session_state['analyzer']['morph_tagger'])
                for token in doc.tokens:
                    token.lemmatize(st.session_state['analyzer']['morph_vocab'])
                chunks = [Substring(token.start, token.stop, token.text) for token in doc.tokens]
                lemmas = [token.lemma for token in doc.tokens]
        else:
            lemmas = [t.text.lower() for t in chunks]
        
        j = 0
        for s in docs[i]['sents']:
            tokens = []
            while j < len(chunks) and chunks[j].start < s.stop:
                tokens.append(Substring(chunks[j].start, chunks[j].stop, lemmas[j]))
                j += 1
            docs[i]['tokens'].append(tokens)
            docs[i]['input'].append(''.join(s.text + ' ' for s in tokens))
        docs[i]['input'] = ''.join(s for s in docs[i]['input'])
        pr_pbar.progress((i+1)/len(docs), text=f'Подготовка словаря: {i+1}/{len(docs)}')
    
    st.session_state['docs'] = docs
    vectorizer = CountVectorizer(token_pattern=token_pattern_filt, stop_words=st.session_state['stopwords'])
    st.session_state['counts'] = vectorizer.fit_transform([d['input'] for d in docs]).toarray()
    st.session_state['words'] = vectorizer.get_feature_names_out()
    
    words = st.session_state['words']
    vocab = {w: i for i,w in enumerate(words)}
    for i in range(len(docs)):
        docs[i]['seq'] = []
        docs[i]['parts'] = []
        for j in range(len(docs[i]['tokens'])):
            docs[i]['tokens'][j] = [t for t in docs[i]['tokens'][j] if t.text in vocab]
            add = [vocab[t.text] for t in docs[i]['tokens'][j]]
            if add:
                docs[i]['seq'] += add
                docs[i]['parts'] += [j for _ in range(len(add))]
        docs[i]['seq'] = np.array(docs[i]['seq'], dtype='int')
        if not byword:
            docs[i]['parts'] = np.array(docs[i]['parts'], dtype='int')
        else:
            docs[i]['parts'] = np.arange(len(docs[i]['parts']), dtype='int')
        docs[i]['tokens'] = sum((t for t in docs[i]['tokens']), [])
    
    st.session_state['ready'] = ''
        
    #with open('words.txt', encoding='utf-8') as f:
    #    words0 = f.read().split('\n')
    #print(set(words) - set(words0))
    
    st.session_state['navec'] = Navec.load('./src/navec_hudlit_v1_12B_500K_300d_100q.tar')
    word_embs = np.full((len(words), 300), np.nan)
    for i,word in enumerate(words):
        if word in st.session_state['navec']:
            word_embs[i] = st.session_state['navec'][word]
    mask = ~np.any(np.isnan(word_embs), axis=1)
    st.session_state['word_sims'] = np.full((len(words), len(words)), np.nan)
    st.session_state['word_sims'][np.ix_(mask,mask)] = 1 - cosine_distances(word_embs[mask])
    
    dist = 10
    min_count = 10
    idxs = np.where(st.session_state['counts'].sum(axis=0) >= min_count)[0]
    
    npmi = np.zeros((len(idxs), len(idxs)))
    probs = {}
    seq = np.hstack([d['seq'] for d in docs])
    p = (2*dist - 1) / max(1, len(seq))
    pos = [np.where(seq == i)[0] for i in idxs]
    for i1 in range(len(idxs)):
        pr_pbar.progress((i1+1)/len(idxs), text=f'Обработка коллокаций: {i1+1}/{len(idxs)}')
        v1 = pos[i1]
        for i2 in range(i1+1,len(idxs)):
            v2 = pos[i2]
            k = min(len(v1),len(v2))
            d = np.abs(v1[...,np.newaxis] - v2) < dist
            if d.any():
                if k > 1:
                    rows,cols = linear_sum_assignment(d, maximize=True)
                    num = d[rows,cols].sum()
                else:
                    num = np.any(d).astype(int)    
                if num > 0:
                    m = max(len(v1),len(v2))
                    if (k,m,num) not in probs:
                        val = 0
                        q = 1 - (1-p) ** m
                        for n in range(num, k+1):
                            val += (q ** n) * ((1-q) ** (k-n)) * math.comb(k, n)                     
                        probs[(k,m,num)] = max(0, 1-val)  
                    npmi[i1,i2] = probs[(k,m,num)]    
                    npmi[i2,i1] = npmi[i1,i2]
    
    st.session_state['word_dist'] = dist
    st.session_state['word_min_count'] = min_count
    st.session_state['npmi'] = npmi

def perplexity_hmm(start, trans, emis, reduce=True):
    use_lengths = False if 'use_lengths' not in st.session_state else st.session_state['use_lengths']
    docs = st.session_state['docs']
    n_topics = len(start)
    perp = 0 if reduce else []
    
    if st.session_state['ready'] == 'SentHMM':
        A = start.copy()
        k = 0
        for doc in docs:
            p = 0
            for j in range(len(doc['sents'])):
                for t in np.where([doc['parts'] == j])[1]:
                    p += np.log(emis[st.session_state['seq'][k], doc['seq'][t]])
                p += np.log(A[st.session_state['seq'][k]])
                A = trans[st.session_state['seq'][k]]
                k += 1
            if use_lengths:
                A = start.copy()
            if reduce:
                perp += p
            else:
                perp.append(p)
                                
    else:
        A = start.copy()
        for doc in docs:
            p = 0
            seq = doc['seq']
            parts = doc['parts']
            for j in range(len(np.unique(doc['parts']))):
                if any(parts == j):
                    if parts[0] != j:
                        A = A @ trans
                    for k in np.where([parts == j])[1]:
                        A *= emis[:,seq[k]]
                        q = A.max()
                        A /= q
                        p += np.log(q)
            p += np.log(A.sum())
            if use_lengths:
                A = start.copy()
            elif len(seq) > 0:
                (A / A.sum()) @ trans
            if reduce:
                perp += p
            else:
                perp.append(p)
                
    if reduce:
        perp = np.exp(-perp / sum(len(d['seq']) for d in docs))
    else:
        perp = np.exp(-np.array(perp) / np.array([max(len(d['seq']),1) for d in docs]))
            
    return perp
    
def perplexity(D, phi, theta):
    temp = np.log(phi @ theta);
    temp[D == 0] = 0;
    p = np.exp((-(D * temp).sum()) / D.sum())
    return p

def correlation(phi):
    p = np.corrcoef(phi.transpose())
    p = (p.sum() - p.shape[0]) / (p.shape[0] * (p.shape[0] - 1))
    return p

def coherence(phi, topk=10, dist=10):
    idx0 = np.argsort(-phi, axis=0)[:topk]
    idx = np.unique(idx0)
    inv_idx = {j: i for i,j in enumerate(idx)}
    occur = np.zeros((len(idx),), dtype='int')
    cooccur = np.zeros((len(idx),len(idx)), dtype='int')
    seqs = [d['seq'] for d in st.session_state['docs'] if len(d['seq']) > 0]
    for seq in seqs:
        ind = np.array([i for i,j in enumerate(idx) if j in seq])
        if len(ind) > 0:
            pos = [np.where(seq == idx[j]) - np.arange(0, max(1, len(seq)-dist+1))[...,np.newaxis] for j in ind]
            pos = np.vstack([np.any((lambda x: np.logical_and(x >= 0, x < dist))(p), axis=1) for p in pos])
            occur[ind] += pos.sum(axis=1)
            cooccur[np.ix_(ind,ind)] += np.minimum(pos[...,np.newaxis], \
                                                   pos.transpose().reshape((1,pos.shape[1],pos.shape[0]))).sum(axis=1)
    ndocs = sum(max(1, len(seq)-dist+1) for seq in seqs)
    c = np.log(cooccur)
    numer = c - np.log(occur * occur[...,np.newaxis] / ndocs)
    denom = np.log(ndocs) - c
    npmi = numer / denom
    npmi[np.isnan(npmi)] = -1
    np.fill_diagonal(npmi, 0)
    
    p = np.zeros((idx0.shape[1],))
    for i in range(idx0.shape[1]):
        ind = np.array([inv_idx[j] for j in idx0[:,i]])
        p[i] = npmi[np.ix_(ind,ind)].sum() / (len(ind) * (len(ind)-1))
    p = p.mean()
    
    return p

def diversity(phi, topk=10):
    idx = np.argsort(-phi, axis=0)[:topk]
    p = len(np.unique(idx)) / idx.size
    return p
    
def fill_result():    
    topicnames = []
    for i in range(st.session_state['phi'].shape[1]):
        indices = np.argsort(st.session_state['phi' if st.session_state['ready'] != 'BERTopic' else 'ctfidf'][:,i])[::-1] 
        temp = f''
        for j in range(7):
            word = st.session_state['words'][indices[j]]
            temp += f'{word}, '
        topicnames.append(temp[:-2])
    st.session_state['topicnames'] = topicnames
    
    if st.session_state['ready'] in ['LDA','ARTM','CTM','BERTopic','Top2Vec']:
        temp = np.log(st.session_state['phi'] @ st.session_state['theta'])
        D = st.session_state['counts'].transpose()
        temp[D == 0] = 0
        p = np.exp((-(D * temp).sum(axis=0)) / D.sum(axis=0))
        p[np.isnan(p)] = 1
        st.session_state['perplexity'] = p
    
    elif st.session_state['ready'] == 'SentHMM':
        n_topics = len(st.session_state['start'])
        theta = np.zeros((n_topics, st.session_state['counts'].shape[0]))
        phi = np.zeros((st.session_state['counts'].shape[1], theta.shape[0]))
        labels = []
        k = 0
        for i,d in enumerate(st.session_state['docs']):
            add = -np.ones_like(d['parts'])
            for j in range(len(d['sents'])):
                add[d['parts'] == j] = st.session_state['seq'][k]
                k += 1
            labels.append(add)
            theta[:,i] += np.bincount(add, minlength=theta.shape[0])
            for a,s in zip(add,d['seq']):
                phi[s,a] += 1
        phi = phi/phi.sum(axis=0)
        phi[np.isnan(phi)] = 1/phi.shape[0]
        theta = theta/theta.sum(axis=0)
        theta[np.isnan(theta)] = 1/theta.shape[0]
        st.session_state['phi'] = phi
        st.session_state['theta'] = theta
        st.session_state['labels'] = labels
        st.session_state['perplexity'] = perplexity_hmm(st.session_state['start'], st.session_state['trans'], \
                                                        st.session_state['phi'].transpose(), reduce=False)
    elif st.session_state['ready'] == 'HMM':
        emis = st.session_state['phi'].transpose()
        n_topics = emis.shape[0]
        theta = np.zeros((n_topics, st.session_state['counts'].shape[0]))
        docs = st.session_state['docs']
        start = st.session_state['start']
        trans = st.session_state['trans']
        st.session_state['perplexity'] = perplexity_hmm(start, trans, emis, reduce=False)
        
        labels = [[] for _ in range(len(docs))]
        for i in range(len(docs)):
            seq = docs[i]['seq']
            k = len(seq)
            if k > 0:
                P = docs[i]['parts'][:-1] != docs[i]['parts'][1:]                
                A = start * emis[:,seq[0]]
                M = np.zeros((n_topics, k-1), dtype=int)
                for j in range(1,k):
                    if P[j-1]:
                        A = A[...,np.newaxis] * trans
                        M[:,j-1] = np.argmax(A, axis=0)
                        A = A.max(axis=0) * emis[:,seq[j]]
                    else:
                        A *= emis[:,seq[j]]
                        M[:,j-1] = np.arange(n_topics)
                    A /= A.max()
                    
                res = [int(np.argmax(A))]
                for j in range(k-2,-1,-1):
                    res = [M[res[0],j]] + res
                labels[i] = np.array(res)
                theta[:,i] = np.bincount(res, minlength=n_topics)
        
        theta = theta / theta.sum(axis=0)
        theta[np.isnan(theta)] = 1/n_topics
        st.session_state['labels'] = labels
        st.session_state['theta'] = theta
        
    colors = [WD_COLOR_INDEX.BLUE, WD_COLOR_INDEX.DARK_BLUE, WD_COLOR_INDEX.DARK_RED, WD_COLOR_INDEX.DARK_YELLOW, WD_COLOR_INDEX.GRAY_50,
                  WD_COLOR_INDEX.GREEN, WD_COLOR_INDEX.PINK, WD_COLOR_INDEX.RED, WD_COLOR_INDEX.TEAL, WD_COLOR_INDEX.VIOLET]
    
    n_topics = st.session_state['phi'].shape[1]
    wdoc = Document()
    for i in range(n_topics):
        wdoc.add_heading(f'Тема {i+1}', level=1)
        p = wdoc.add_paragraph()
        indices = np.argsort(st.session_state['phi'][:,i])[::-1]
        for j in range(20):
            word = st.session_state['words'][indices[j]]
            score = st.session_state['phi'][indices[j],i]
            r = p.add_run(f'{word} ({score:.4f}) ')
            r.font.color.rgb = RGBColor(255,255,255)
            r.font.highlight_color = colors[i]

    if st.session_state['ready'] in ['HMM', 'SentHMM']: 
        font_size = 10
        wdoc.add_heading(f'Таблица переходов', level=1)
        tab = wdoc.add_table(n_topics+2, n_topics+1)
        tab.style = 'Table Grid'
        cell = tab.cell(0, 0)
        par = cell.paragraphs[0]
        par.alignment = WD_ALIGN_PARAGRAPH.CENTER
        cell = tab.cell(1, 0)
        par = cell.paragraphs[0]
        par.alignment = WD_ALIGN_PARAGRAPH.CENTER
        run = par.add_run(f'0')
        run.bold = True
        run.font.size = Pt(font_size)
        for j in range(n_topics):
            cell = tab.cell(0, j+1)
            par = cell.paragraphs[0]
            par.alignment = WD_ALIGN_PARAGRAPH.CENTER
            val = st.session_state['trans'][i,j]
            run = par.add_run(f'{j+1}')
            run.bold = True
            run.font.size = Pt(font_size)
            cell = tab.cell(1, j+1)
            par = cell.paragraphs[0]
            par.alignment = WD_ALIGN_PARAGRAPH.CENTER
            val = st.session_state['start'][j]
            run = par.add_run(f'{val:.4f}')
            run.font.size = Pt(font_size)
        for i in range(n_topics):
            cell = tab.cell(i+2, 0)
            par = cell.paragraphs[0]
            par.alignment = WD_ALIGN_PARAGRAPH.CENTER
            run = par.add_run(f'{i+1}')
            run.bold = True
            run.font.size = Pt(font_size)
            for j in range(n_topics):
                cell = tab.cell(i+2, j+1)
                par = cell.paragraphs[0]
                par.alignment = WD_ALIGN_PARAGRAPH.CENTER
                val = st.session_state['trans'][i,j]
                run = par.add_run(f'{val:.4f}')
                run.font.size = Pt(font_size)

    if st.session_state['ready'] in markup_cases:
        docs = st.session_state['docs']
        for i in range(len(docs)):
            wdoc.add_heading(docs[i]['name'], level=1)
            p = wdoc.add_paragraph()

            if st.session_state['ready'] not in ['HMM','SentHMM']:
                prev,k0 = 0,-1
                for j,pos in enumerate(st.session_state['docs'][i]['seq']):
                    tok = st.session_state['docs'][i]['tokens'][j]
                    r0 = p.add_run(st.session_state['docs'][i]['text'][prev:tok.start])
                    temp = st.session_state['phi'][pos,:] * st.session_state['theta'][:,i]
                    k = np.argmax(temp) % len(rgb)
                    r = p.add_run(st.session_state['docs'][i]['text'][tok.start:tok.stop])
                    if (temp / temp.sum()).max() > 5/n_topics:
                        if k == k0:
                            r0.font.color.rgb = RGBColor(255,255,255)
                            r0.font.highlight_color = colors[k]
                        elif k0 >= 0 and not any(c.isalpha() for c in r0.text):
                            r0.font.color.rgb = RGBColor(255,255,255)
                            r0.font.highlight_color = colors[k0]
                        r.font.color.rgb = RGBColor(255,255,255)
                        r.font.highlight_color = colors[k]
                        k0 = k
                    else:
                        k0 = -1
                    prev = tok.stop
                r = p.add_run(st.session_state['docs'][i]['text'][prev:])
            else:
                prev = 0
                k = 0 if len(st.session_state['labels'][i]) == 0 else st.session_state['labels'][i][0]
                for j,sent in enumerate(st.session_state['docs'][i]['sents']):
                    r = p.add_run(st.session_state['docs'][i]['text'][prev:sent.start])
                    r.font.color.rgb = RGBColor(255,255,255)
                    r.font.highlight_color = colors[k]
                    if any(st.session_state['docs'][i]['parts'] == j):
                        k = mode(st.session_state['labels'][i][st.session_state['docs'][i]['parts'] == j]).mode[0] % len(colors)
                    r = p.add_run(st.session_state['docs'][i]['text'][sent.start:sent.stop])
                    r.font.color.rgb = RGBColor(255,255,255)
                    r.font.highlight_color = colors[k]
                    prev = sent.stop
                r = p.add_run(st.session_state['docs'][i]['text'][prev:])
                r.font.color.rgb = RGBColor(255,255,255)
                r.font.highlight_color = colors[k]

        wdoc.save('Разметка.docx')
        with open('Разметка.docx','rb') as f:
            st.session_state['markup'] = f.read()
    
    nwt = st.session_state['phi'] * (st.session_state['theta'] * st.session_state['counts'].sum(axis=1)).sum(axis=1)
    df = pd.DataFrame(nwt / nwt.sum(axis=1)[...,np.newaxis])
    df.index = st.session_state['words']
    st.session_state['phi_df'] = df
    st.session_state['phi_df'].columns = np.arange(len(st.session_state['phi_df'].columns))+1
    
    df = pd.DataFrame(st.session_state['theta'].transpose())
    df.index = [d['name'] for d in st.session_state['docs']]
    st.session_state['theta_df'] = df
    st.session_state['theta_df'].columns = np.arange(len(st.session_state['theta_df'].columns))+1
    
    
def init_phi_theta(n_topics, w = 1, random_state=None):
    counts = st.session_state['counts']

    mask = np.any(counts, axis=1)
    X = counts[mask] / counts[mask].sum(axis=1)[...,np.newaxis]
    labels = np.zeros((counts.shape[0],), dtype=int)
    kmeans = SphericalKMeans(n_clusters=n_topics, random_state=random_state).fit(csr_matrix(X))
    labels[mask] = kmeans.labels_
    
    if byword:
        labels[mask] = np.load('labels.npy')[0]

    sub0 = np.argsort([counts.sum(axis=1)[labels == i].sum() for i in range(n_topics)])
    sub = sub0.copy()
    sub[sub0] = np.arange(n_topics)
    labels = sub[labels]
    
    theta = w * np.ones((n_topics, counts.shape[0]))
    for i in range(theta.shape[1]):
        theta[labels[i],i] += 1

    phi = (theta @ counts).transpose()
    phi = phi / phi.sum(axis=0)
    theta = theta / theta.sum(axis=0)
    
    return phi, theta

def calc_metrics(counts, phi, theta):
    corr = correlation(phi)
    coh = coherence(phi)
    div = diversity(phi)
    return ({'Перплексия': perplexity(counts, phi, theta)} if theta is not None else {}) | \
        {'Корреляция': corr, 'Когерентность': coh, 'Разнообразие': div}


def main():
    
    st.title('Тематическое моделирование')              
    
    if 'ready' not in st.session_state:
        st.session_state['ready'] = ''
    if 'metrics' not in st.session_state:
        st.session_state['metrics'] = {}
        
    with st.sidebar:
        doc_fn = st.file_uploader('Загрузить документы', type='txt')
        stop_fn = st.file_uploader('Загрузить стоп-слова', type='txt')
        norm_cb = st.checkbox('Перевести в современную орфографию', value=True)
        lemma_cb = st.checkbox('Сделать лемматизацию', value=False)
        vocab_btn = st.button('Получить словарь')
        pr_pbar = st.progress(0.0, text='Подготовка словаря:')
        random_state = st.number_input('Cлучайное зерно', step=1)
        
    if doc_fn:
        if 'doc_id' not in st.session_state or st.session_state['doc_id'] != doc_fn.file_id:
            st.session_state['doc_id'] = doc_fn.file_id
            docs = StringIO(doc_fn.getvalue().decode('utf-8'))
            docs = docs.read()
            docs = [d.strip('\r').split('\t') for d in docs.split('\n')]
            docs = [{'name': d[0], 'date': d[1], 'text_orig': d[2]} for d in docs if d[0]]
            st.session_state['docs'] = docs
            st.session_state['ready'] = ''
            st.session_state['words'] = []
        else:
            docs = st.session_state['docs']
    else:
        docs,st.session_state['docs'],st.session_state['words'] = [],[],[]
    
    if stop_fn:
        if 'stop_id' not in st.session_state or st.session_state['stop_id'] != stop_fn.file_id:
            st.session_state['stop_id'] = stop_fn.file_id
            stopwords = StringIO(stop_fn.getvalue().decode('utf-8'))
            stopwords = stopwords.read().split()
            st.session_state['stopwords'] = stopwords
            st.session_state['ready'] = ''
        else:
            stopwords = st.session_state['stopwords']
    else:
        stopwords,st.session_state['stopwords'] = [],[]
    
    if vocab_btn and docs:
        prepare_vocab(norm_cb, lemma_cb, pr_pbar)
    
    col1, col2 = st.columns([0.33, 0.67])
    with col1:
        tab1,tab2,tab3,tab4,tab5,tab6 = st.tabs(['LDA','ARTM','HMM','CTM','BERTopic','Top2Vec'])
        
        with tab1: # LDA
            n_topics = st.slider('Число тем', min_value=2, max_value=100, value=10)
            theta_reg = st.slider('Коэффициент сглаживания тем в документах', min_value=0.01, max_value=1.0, value=0.1, step=0.01)
            phi_reg = st.slider('Коэффициент сглаживания слов в темах', min_value=0.01, max_value=1.0, value=0.1, step=0.01)
            max_iter = st.slider('Число итераций', min_value=1, max_value=1000, value=100)
            empty_msg = 'Извлечение тем  \nПолная перплексия  \nПерплексия  \nКорреляция  \nКогерентность  \nРазнообразие'
            if st.button('Извлечь темы') and docs:
                tm_pbar = st.progress(0.0, text=empty_msg)
                if 'words' not in st.session_state or len(st.session_state['words']) == 0:
                    prepare_vocab(norm_cb, lemma_cb, pr_pbar)
                lda = LatentDirichletAllocation(n_components=n_topics, doc_topic_prior=theta_reg, topic_word_prior=phi_reg, \
                                                max_iter=max_iter, evaluate_every=1, perp_tol=0.01, verbose=True, random_state=random_state)
                sys.stdout_old = sys.stdout
                sys.stdout = StringIO()
                t = Thread(target=lda.fit_transform, args=(st.session_state['counts'],))
                t.start()
                while t.is_alive():
                    time.sleep(0.1)
                    text = sys.stdout.getvalue()
                    m = re.findall('iteration: ([0-9]+) of max_iter: ([0-9]+), perplexity: ([0-9]+.?[0-9]+)', text)                   
                    if m:
                        phi = lda.components_.transpose()
                        phi = phi / phi.sum(axis=0)
                        theta = lda.transform(st.session_state['counts']).transpose()
                        metrics = {'Полная перплексия': float(m[-1][2])} | calc_metrics(st.session_state['counts'].transpose(), phi, theta)
                        st.session_state['progress_msg'] = f'Извлечение тем: {m[-1][0]}/{m[-1][1]}  \n' + \
                            ''.join(f'{k}: {v:.4f}  \n' for k,v in metrics.items())
                        tm_pbar.progress(int(m[-1][0])/int(m[-1][1]), text=st.session_state['progress_msg'])
                t.join()
                text = sys.stdout.getvalue()
                sys.stdout = sys.stdout_old
                
                phi = lda.components_.transpose()
                phi = phi / phi.sum(axis=0)
                theta = lda.transform(st.session_state['counts']).transpose()
                m = re.findall('iteration: ([0-9]+) of max_iter: ([0-9]+), perplexity: ([0-9]+.?[0-9]+)', text)
                metrics = {'Полная перплексия': float(m[-1][2]) if m else 1.0}
                metrics = metrics | calc_metrics(st.session_state['counts'].transpose(), phi, theta)
                st.session_state['progress_msg'] = f'Извлечение тем: {m[-1][0] if m else max_iter}/{max_iter}  \n' + \
                    ''.join(f'{k}: {v:.4f}  \n' for k,v in metrics.items())
                tm_pbar.progress(1.0, text=st.session_state['progress_msg'])
                
                #theta_emb = PCA(n_components=2).fit_transform(st.session_state['counts'])
                theta_emb = UMAP(n_components=2).fit_transform(st.session_state['counts'])
                #theta_emb = TSNE(n_components=2).fit_transform(st.session_state['counts'])
                st.session_state['theta_emb'] = theta_emb
                
                st.session_state['phi'] = phi
                st.session_state['theta'] = theta
                st.session_state['ready'] = 'LDA'
                fill_result()
            else:
                tm_pbar = st.progress(1.0 if st.session_state['ready'] == 'LDA' else 0.0, \
                                      text=st.session_state['progress_msg'] if st.session_state['ready'] == 'LDA' else empty_msg)
                
        with tab2: # ARTM
            n_topics = st.slider('Число тем', min_value=2, max_value=100, value=10, key='n_topics2')
            theta_reg = st.slider('Коэффициент сглаживания тем в документах', min_value=-10.0, max_value=10.0, value=0.0, step=0.01, \
                                  key='theta_reg2')
            phi_reg = st.slider('Коэффициент сглаживания слов в темах',min_value=-10.0, max_value=10.0, value=0.0, step=0.01, key='phi_reg2')
            decorr_reg = st.slider('Коэффициент контрастирования тем',min_value=0, max_value=100000, value=0, step=100, key='decorr_reg2')
            emb_reg = st.slider('Коэффициент сглаживания по синонимам',min_value=0, max_value=1000, value=0, step=1, key='emb_reg2')
            max_iter = st.slider('Число итераций', min_value=1, max_value=1000, value=100, key='max_iter2')
            empty_msg = 'Извлечение тем  \nПерплексия  \nКорреляция  \nКогерентность  \nРазнообразие'
                        
            if st.button('Извлечь темы', key='run_tm2') and docs:
                tm_pbar2 = st.progress(0.0, text=empty_msg)
                if 'words' not in st.session_state or len(st.session_state['words']) == 0:
                    prepare_vocab(norm_cb, lemma_cb, pr_pbar)
                
                word_sims = st.session_state['word_sims']
                word_sims = np.nan_to_num(word_sims, nan=0.0)
                word_sims -= np.diag(word_sims.sum(axis=1))
                
                phi,theta = init_phi_theta(n_topics, w=1, random_state=random_state)
                metrics = calc_metrics(st.session_state['counts'].transpose(), phi, theta)
                st.session_state['progress_msg'] = f'Извлечение тем: 0/{max_iter}  \n' + \
                    ''.join(f'{k}: {v:.4f}  \n' for k,v in metrics.items())
                tm_pbar2.progress(0.0, text=st.session_state['progress_msg'])

                part = 100
                for i in range(max_iter):
                    
                    phi0 = np.zeros_like(phi)
                    theta0 = np.zeros_like(theta)
                    for a in range((phi.shape[0]-1) // part + 1):
                        it = [a*part, min((a+1)*part, phi.shape[0])]
                        phi1 = phi[it[0]:it[1],:]
                        for b in range((theta.shape[1]-1) // part + 1):
                            jt = [b*part, min((b+1)*part, theta.shape[1])]
                            theta1 = theta[:,jt[0]:jt[1]]
                            ntdw = phi1[...,np.newaxis] * theta1.reshape((1,) + theta1.shape)
                            ntdw = ntdw / ntdw.sum(axis=1)[:,np.newaxis,:]
                            ntdw[np.isnan(ntdw)] = 1/n_topics
                            ntdw = st.session_state['counts'][jt[0]:jt[1],it[0]:it[1]].transpose()[:,np.newaxis,:] * ntdw
                            phi0[it[0]:it[1],:] += ntdw.sum(axis=2)
                            theta0[:,jt[0]:jt[1]] += ntdw.sum(axis=0)
                    
                    #ntdw = phi[...,np.newaxis] * theta.reshape((1,) + theta.shape)
                    #ntdw = ntdw / ntdw.sum(axis=1)[:,np.newaxis,:]
                    #ntdw[np.isnan(ntdw)] = 1/n_topics
                    #ntdw = st.session_state['counts'].transpose()[:,np.newaxis,:] * ntdw
                    
                    phi = phi0 + phi_reg
                    if decorr_reg > 0 or emb_reg > 0:
                        phi0 = phi0 / phi0.sum(axis=0)
                        phi0[np.isnan(phi0)] = 1/phi0.shape[0]
                    if decorr_reg > 0:
                        phi -= decorr_reg*phi0*(phi0.sum(axis=1)[...,np.newaxis]-phi0)
                    if emb_reg > 0:
                        #phi = (emb_reg * word_sims @ phi0 + phi0) / (emb_reg * word_sims.sum(axis=1)[...,np.newaxis] + 1)
                        phi += emb_reg * phi0 * (word_sims @ phi0)
                    
                    phi = np.maximum(phi, 0)
                    phi = phi / phi.sum(axis=0)
                    phi[np.isnan(phi)] = 1/phi.shape[0]

                    theta = np.maximum(theta0 + theta_reg, 0.0)
                    theta = theta / theta.sum(axis=0)
                    theta[np.isnan(theta)] = 1/theta.shape[0]
                    
                    metrics = calc_metrics(st.session_state['counts'].transpose(), phi, theta)
                    st.session_state['progress_msg'] = f'Извлечение тем: {i+1}/{max_iter}  \n' + \
                        ''.join(f'{k}: {v:.4f}  \n' for k,v in metrics.items())
                    tm_pbar2.progress((i+1)/max_iter, text=st.session_state['progress_msg'])
                    
                st.session_state['theta_emb'] = UMAP(n_components=2).fit_transform(st.session_state['counts'])
                st.session_state['phi'] = phi
                st.session_state['theta'] = theta
                st.session_state['ready'] = 'ARTM'
                fill_result()
            else:
                tm_pbar2 = st.progress(1.0 if st.session_state['ready'] == 'ARTM' else 0.0, \
                                       text=st.session_state['progress_msg'] if st.session_state['ready'] == 'ARTM' else empty_msg)
            
        with tab3: # HMM
            n_topics = st.slider('Число тем', min_value=2, max_value=100, value=10, key='n_topics3')
            theta_reg = st.slider('Коэффициент сглаживания тем в документах', min_value=-10.0, max_value=10.0, value=0.0, step=0.01, \
                                  key='theta_reg3')
            phi_reg = st.slider('Коэффициент сглаживания слов в темах', min_value=-10.0, max_value=10.0, value=0.0, step=0.01, \
                                key='phi_reg3')
            mean_len = st.slider('Средняя длина темы', min_value=2, max_value=100, value=10)
            max_iter = st.slider('Число итераций', min_value=0, max_value=1000, value=100, key='max_iter3')
            use_prev = st.checkbox('Стартовать с предыдущего результата', value=True, key='use_prev3')
            empty_msg = 'Извлечение тем  \nПерплексия  \nКорреляция  \nКогерентность  \nРазнообразие'
            
            if st.button('Извлечь темы', key='run_tm3') and docs:
                tm_pbar3 = st.progress(0.0, text=empty_msg)
                if 'words' not in st.session_state or len(st.session_state['words']) == 0:
                    prepare_vocab(norm_cb, lemma_cb, pr_pbar)
                
                if st.session_state['ready'] in ['LDA','ARTM','HMM'] and use_prev:
                    phi,theta = st.session_state['phi'],st.session_state['theta']
                else:
                    phi,theta = init_phi_theta(n_topics, w=1, random_state=random_state)
                    
                emis = phi.transpose()
                if st.session_state['ready'] == 'HMM' and use_prev:
                    start = st.session_state['start']
                    trans = st.session_state['trans']
                else:
                    start = np.ones((n_topics, )) / n_topics
                    trans = np.ones((n_topics, n_topics)) + ((mean_len-1)*(n_topics-1)-1)*np.eye(n_topics)
                    trans = trans / trans.sum(axis=1)
                
                use_lib = False
                st.session_state['use_lengths'] = True
                if use_lib:
                    model = CategoricalHMM(n_components=n_topics, verbose=True, n_iter=max_iter, algorithm='map', init_params='')
                    model.startprob_ = start
                    model.transmat_ = trans
                    model.emissionprob_ = emis
                    model.startprob_prior = theta_reg + 1
                    model.transmat_prior = theta_reg + 1
                    model.emissionprob_prior = phi_reg + 1
                    
                    sys.stderr_old = sys.stderr
                    sys.stderr = StringIO()
                    t = Thread(target=model.fit, args=(np.hstack(d['seq'] for d in docs).reshape(-1,1), \
                                                       [len(d['seq']) for d in docs if len(d['seq']) > 0]))
                    t.start()
                    while t.is_alive():
                        time.sleep(0.1)
                        text = sys.stderr.getvalue()
                        m,r = None,iter(text.split('\n')[::-1])
                        while (s := next(r, None)) is not None and not m:
                            m0 = re.findall('([0-9]+)[ ]+([+-]?[0-9]+.?[0-9]+)', s)
                            if m0:
                                m = m0
                        if m:
                            metrics = {'Полная перплексия': -float(m[-1][1])}
                            metrics['Перплексия'] = perplexity_hmm(model.startprob_, model.transmat_, model.emissionprob_)
                            metrics = metrics | calc_metrics(st.session_state['counts'].transpose(), model.emissionprob_.transpose(), None)
                            st.session_state['progress_msg'] = f'Извлечение тем: {m[-1][0]}/{max_iter}  \n' + \
                            ''.join(f'{k}: {v:.4f}  \n' for k,v in metrics.items())
                            tm_pbar3.progress(int(m[-1][0])/max_iter, text=st.session_state['progress_msg'])
                    t.join()
                    text = sys.stderr.getvalue()
                    sys.stderr = sys.stderr_old
                    
                    m,r = None,iter(text.split('\n')[::-1])
                    while (s := next(r, None)) is not None and not m:
                        m0 = re.findall('([0-9]+) ([+-]?[0-9]+.?[0-9]+)', s)
                        if m0:
                            m = m0
                    
                    metrics = {'Полная перплексия': -float(m[-1][1]) if m else 1}
                    metrics['Перплексия'] = perplexity_hmm(model.startprob_, model.transmat_, model.emissionprob_)
                    metrics = metrics | calc_metrics(st.session_state['counts'].transpose(), model.emissionprob_.transpose(), None)
                    
                    emis = model.emissionprob_
                    start = model.startprob_
                    trans = model.transmat_
                    
                    st.session_state['progress_msg'] = f'Извлечение тем: {max_iter}/{max_iter}  \n' + \
                        ''.join(f'{k}: {v:.4f}  \n' for k,v in metrics.items())
                    tm_pbar3.progress(1.0, text=st.session_state['progress_msg'])
                    
                for i in range(max_iter if not use_lib else 0):
                    start0 = np.zeros_like(start)
                    trans0 = np.zeros_like(trans)
                    emis0 = np.zeros_like(emis)
                    p = 0
                    for pos in range(len(docs)):
                        seq = docs[pos]['seq']
                        if seq.size > 0:
                            
                            P = docs[pos]['parts'][:-1] != docs[pos]['parts'][1:]
                            k = len(seq)
                            V = np.zeros((k, emis.shape[1]))
                            V[np.arange(k), seq] = 1 

                            A = np.zeros((n_topics, k))
                            B = np.zeros((n_topics, k))
                            H = np.zeros((n_topics, n_topics, k))                            
                            
                            for j in range(k):
                                if j == 0:
                                    A[:,0] = start * emis[:,seq[0]]
                                elif P[j-1]:
                                    A[:,j] = (A[:,j-1] @ trans) * emis[:,seq[j]]
                                else:
                                    A[:,j] = A[:,j-1] * emis[:,seq[j]]
                                q = A[:,j].max()
                                A[:,j] /= q
                                p += np.log(q)
                            p += np.log(A[:,-1].sum())
                            
                            B[:,-1] = 1
                            for j in range(k-2,-1,-1):
                                if P[j]:
                                    B[:,j] = trans @ (B[:,j+1] * emis[:,seq[j+1]])
                                else:
                                    B[:,j] = B[:,j+1] * emis[:,seq[j+1]]
                                B[:,j] /= B[:,j].max()
                                
                            G = A * B;
                            G /= G.sum(axis=0)
                            for j in range(k-1):
                                if P[j]:
                                    H[:,:,j] = (A[:,j:j+1] @ B[:,j+1:j+2].transpose()) * trans * emis[:,seq[j+1]:seq[j+1]+1].transpose()
                                    H[:,:,j] /= H[:,:,j].sum()
                                    
                            start0 += G[:,0]
                            trans0 += H.sum(axis=2)
                            emis0 += G @ V
                    
                    start0 += theta_reg
                    trans0 += theta_reg
                    start = start0 / start0.sum()
                    trans = trans0 / trans0.sum(axis=1)[...,np.newaxis]
                    emis = np.maximum(0, emis0 + phi_reg)
                    emis /= emis.sum(axis=1)[...,np.newaxis]
                    emis[np.isnan(emis)] = 1 / emis.shape[1]
                    
                    metrics = {'Перплексия': np.exp(-p / sum([len(d['seq']) for d in docs]))}
                    metrics = metrics | calc_metrics(st.session_state['counts'].transpose(), emis.transpose(), theta=None)
                    st.session_state['progress_msg'] = f'Извлечение тем: {i+1}/{max_iter}  \n' + \
                        ''.join(f'{k}: {v:.4f}  \n' for k,v in metrics.items())
                    tm_pbar3.progress((i+1)/max_iter, text=st.session_state['progress_msg'])
                    
                st.session_state['theta_emb'] = UMAP(n_components=2).fit_transform(st.session_state['counts'])
                st.session_state['phi'] = emis.transpose()
                st.session_state['start'] = start
                st.session_state['trans'] = trans
                st.session_state['ready'] = 'HMM'
                fill_result()                
            else:
                tm_pbar3 = st.progress(1.0 if st.session_state['ready'] == 'HMM' else 0.0, \
                                       text=st.session_state['progress_msg'] if st.session_state['ready'] == 'HMM' else empty_msg)
        
        with tab4: # CTM
            n_topics = st.slider('Число тем', min_value=2, max_value=100, value=10, key='n_topics4')
            max_iter = st.slider('Число итераций', min_value=1, max_value=1000, value=100, key='max_iter4')
            
            model_names = ['deeppavlov/rubert-base-cased-sentence',
                           'ai-forever/sbert_large_mt_nlu_ru',
                           'ai-forever/sbert_large_nlu_ru',
                           'ai-forever/ru-en-RoSBERTa',
                           'ai-forever/FRIDA'];
            model_name = st.selectbox("Выберите модель", model_names, key='model_name4')
            qt = TopicModelDataPreparation(model_name, max_seq_length = 512)
            empty_msg = 'Извлечение тем  \nОшибка  \nПерплексия  \nКорреляция  \nКогерентность  \nРазнообразие'
            if st.button('Извлечь темы', key='run_tm4') and docs:
                if 'words' not in st.session_state or len(st.session_state['words']) == 0:
                    prepare_vocab(norm_cb, lemma_cb, pr_pbar)
                
                tm_pbar4 = st.progress(0.0, text='Подготовка данных')
                #training_dataset = qt.fit(text_for_contextual=[d['text'] for d in docs], \
                #                          text_for_bow=[d['input'] for d in docs])
                sys.stderr_old = sys.stderr
                sys.stderr = StringIO()
                
                vectorizer = CountVectorizer(token_pattern=token_pattern_filt, stop_words=st.session_state['stopwords'])
                t = ReturnableThread(target=tm_data_prep_fit, args=(qt, [d['text'] for d in docs], [d['input'] for d in docs], vectorizer))
                t.start()
                while t.is_alive():
                    time.sleep(0.1)
                    text = sys.stderr.getvalue()
                    m = re.findall('([0-9]+)/([0-9]+)', text)
                    if m:
                        tm_pbar4.progress(min(1.0, (int(m[-1][0])+1)/int(m[-1][1])), text=f'Подготовка данных: {int(m[-1][0])+1}/{m[-1][1]}')
                t.join()
                training_dataset = t.result
                if m:
                    tm_pbar4.progress(1.0, text=f'Подготовка данных: {int(m[-1][1])+1}/{m[-1][1]}')
                
                tm_pbar4.progress(0.0, text=empty_msg)
                
                st.session_state['topic_model'] = CombinedTM(bow_size=len(qt.vocab), contextual_size=768, n_components=n_topics, \
                                                             num_data_loader_workers=0, num_epochs=max_iter)
                sys.stderr_old = sys.stderr
                sys.stderr = StringIO()
                
                t = Thread(target=st.session_state['topic_model'].fit, args=(training_dataset,))               
                t.start()
                while t.is_alive():
                    time.sleep(0.1)
                    text = sys.stderr.getvalue()
                    m,r = None,iter(text.split('\r')[::-1])
                    while (s := next(r, None)) is not None and not m:
                        m0 = re.findall('Epoch: \[([0-9]+)/([0-9]+)\].*Train Loss: ([0-9]+.?[0-9]+)', s)
                        if m0:
                            m = m0
                    if m:
                        #lda_vis_data = st.session_state['topic_model'].get_ldavis_data_format(qt.vocab, training_dataset, n_samples=10000)
                        #theta = lda_vis_data['doc_topic_dists'].transpose()
                        #phi = torch.softmax(20*st.session_state['topic_model'].best_components, axis=1).detach().cpu().numpy().transpose()
                        metrics = {'Ошибка': float(m[-1][2])} #| calc_metrics(st.session_state['counts'].transpose(), phi, theta)
                        st.session_state['progress_msg'] = f'Извлечение тем: {m[-1][0]}/{m[-1][1]}  \n' + \
                            ''.join(f'{k}: {v:.4f}  \n' for k,v in metrics.items())
                        tm_pbar4.progress(int(m[-1][0])/int(m[-1][1]), text=st.session_state['progress_msg'])
                t.join()
                
                st.session_state['lda_vis_data'] = st.session_state['topic_model'].get_ldavis_data_format(qt.vocab, training_dataset, \
                                                                                                          n_samples=10000)
                st.session_state['theta'] = st.session_state['lda_vis_data']['doc_topic_dists'].transpose()
                st.session_state['phi'] = torch.softmax(20*st.session_state['topic_model'].best_components, \
                                                        axis=1).detach().cpu().numpy().transpose()
                #phi = (st.session_state['theta'] @ st.session_state['counts']).transpose()
                
                text = sys.stderr.getvalue()
                sys.stderr = sys.stderr_old
                m,r = None,iter(text.split('\r')[::-1])
                while (s := next(r, None)) is not None and not m:
                    m0 = re.findall('Epoch: \[([0-9]+)/([0-9]+)\].*Train Loss: ([0-9]+.?[0-9]+)', s)
                    if m0:
                        m = m0
                metrics = {'Ошибка': float(m[-1][2]) if m else 0.0} | calc_metrics(st.session_state['counts'].transpose(), \
                                                                                   st.session_state['phi'], st.session_state['theta'])
                st.session_state['progress_msg'] = f'Извлечение тем: {m[-1][0] if m else max_iter}/{max_iter}  \n' + \
                    ''.join(f'{k}: {v:.4f}  \n' for k,v in metrics.items())
                tm_pbar4.progress(1.0, text=st.session_state['progress_msg'])
                
                st.session_state['ready'] = 'CTM'
                fill_result()
                
            else:
                tm_pbar4 = st.progress(1.0 if st.session_state['ready'] == 'CTM' else 0.0, \
                                       text=st.session_state['progress_msg'] if st.session_state['ready'] == 'CTM' else empty_msg)
            
        with tab5: # BERTopic
            n_topics = st.slider('Число тем', min_value=2, max_value=100, value=10, key='n_topics5')
            n_mix = st.slider('Число гауссиан', min_value=1, max_value=10, value=3, key='n_mix5')
            mean_len = st.slider('Средняя длина темы', min_value=2, max_value=100, value=5, key='mean_len5')
            max_iter = st.slider('Число итераций', min_value=0, max_value=1000, value=0, key='max_iter5')
            model_names = ['deeppavlov/rubert-base-cased-sentence',
                           'ai-forever/sbert_large_mt_nlu_ru',
                           'ai-forever/sbert_large_nlu_ru',
                           'ai-forever/ru-en-RoSBERTa',
                           'ai-forever/FRIDA'];
            model_name = st.selectbox("Выберите модель", model_names, key='model_name5')
            empty_msg = 'Извлечение тем  \nПерплексия  \nКорреляция  \nКогерентность  \nРазнообразие'
            if st.button('Извлечь темы', key='run_tm5') and docs:
                tm_pbar5 = st.progress(0.0, text=empty_msg)
                
                if 'words' not in st.session_state or len(st.session_state['words']) == 0:
                    prepare_vocab(norm_cb, lemma_cb, pr_pbar)
                
                tm_pbar5.progress(1/4, text=f'Загрузка модели: 1/4')
                embedding_model = SentenceTransformer(model_name)
                try:
                    max_length = embedding_model[0].auto_model.embeddings.position_embeddings.num_embeddings
                    embedding_model[0].max_seq_length = max_length
                    embedding_model.tokenizer.model_max_length = max_length
                except:
                    pass
                vectorizer_model = CountVectorizer(stop_words=stopwords, token_pattern=token_pattern_filt)
                topic_model = BERTopic(embedding_model=embedding_model, vectorizer_model=vectorizer_model, nr_topics=n_topics, \
                                       min_topic_size=2 if max_iter > 0 else n_mix)
                topic_model.embedding_model = select_backend(topic_model.embedding_model, language=topic_model.language, \
                                                             verbose=topic_model.verbose)
                topic_model.umap_model.random_state = random_state
                st.session_state['topic_model'] = topic_model

                tm_pbar5.progress(2/4, text=f'Вычисление эмбеддингов документов: 2/4')
                embeddings = st.session_state['topic_model']._extract_embeddings([d['text'] for d in docs], method="document")
                st.session_state['embeddings'] = embeddings
                tm_pbar5.progress(3/4, text=f'Понижение размерности эмбеддингов: 3/4')
                umap_model = UMAP(n_neighbors=10, n_components=2, min_dist=0.0, metric="cosine", random_state=random_state)
                umap_model.fit(embeddings)
                st.session_state['reduced_embeddings'] = umap_model.embedding_
                tm_pbar5.progress(4/4, text=f'Извлечение тем: 4/4')

                topics, probs = st.session_state['topic_model'].fit_transform([d['text'] for d in docs], embeddings=embeddings)
                n_topics = np.unique(topics).size
                
                st.session_state['use_lengths'] = False
                if max_iter > 0:
                    if False:
                        tm_pbar5.progress(0.0, text=f'Вычисление эмбеддингов предложений: 0/{max_iter}')
                        embs = []
                        for i in range(len(docs)):
                            embs.append(st.session_state['topic_model'].embedding_model.embed([s.text for s in docs[i]['sents']]))
                            tm_pbar5.progress((i+1)/len(docs), text=f'Вычисление эмбеддингов предложений: {i+1}/{len(docs)}')
                        embs = np.vstack(embs)
                        tm_pbar5.progress(1/2, text=f'Понижение размерности эмбеддингов: 1/2')
                        st.session_state['embeddings_sent'] = st.session_state['topic_model'].umap_model.transform(embs)
                        st.session_state['reduced_embeddings_sent'] = umap_model.transform(embs)
                    
                    embs = st.session_state['embeddings_sent']
                    lengths = [len(d['sents']) for d in docs]
                    idxs = np.hstack([(t+1)*np.ones(len(d['sents']),dtype='int') for d,t in \
                                      zip(docs,st.session_state['topic_model'].topics_)])
                    
                    tm_pbar5.progress(2/2, text=f'Вычисление начального приближения гауссиан: 2/2')
                    gmms = [GaussianMixture(n_mix).fit(embs[idxs == i]) for i in range(idxs.max()+1)]
                    weights = np.stack([gmm.weights_ for gmm in gmms])
                    means = np.stack([gmm.means_ for gmm in gmms])
                    covars = np.stack([gmm.covariances_ for gmm in gmms])
                    start = np.ones((n_topics, )) / n_topics
                    trans = np.ones((n_topics, n_topics)) + ((mean_len-1)*(n_topics-1)-1)*np.eye(n_topics)
                    trans = trans / trans.sum(axis=1)
                    
                    use_lib = False
                    if use_lib:
                        gmmhmm = GMMHMM(n_components=phi.shape[1], n_mix=n_mix, covariance_type='full', algorithm='map', \
                                        n_iter=max_iter, verbose=True, init_params='', random_state=random_state)
                        gmmhmm.weights_ = weights
                        gmmhmm.means_ = means
                        gmmhmm.startprob_ = start
                        gmmhmm.transmat_ = trans

                        sys.stderr_old = sys.stderr
                        sys.stderr = StringIO()
                        t = Thread(target=gmmhmm.fit, args=(embs, lengths if st.session_state['use_lengths'] else None))
                        t.start()
                        while t.is_alive():
                            time.sleep(0.1)
                            text = sys.stderr.getvalue()
                            m,r = None,iter(text.split('\n')[::-1])
                            while (s := next(r, None)) is not None and not m:
                                m0 = re.findall('([0-9]+)[ ]+([+-]?[0-9]+.?[0-9]+)', s)
                                if m0:
                                    m = m0
                            if m:
                                metrics = {'Полная перплексия': -float(m[-1][1])}
                                st.session_state['progress_msg'] = f'Подгонка кластеров: {m[-1][0]}/{max_iter}  \n' + \
                                    ''.join(f'{k}: {v:.4f}  \n' for k,v in metrics.items())
                                tm_pbar5.progress(int(m[-1][0])/max_iter, text=st.session_state['progress_msg'])
                        t.join()
                        text = sys.stderr.getvalue()
                        sys.stderr = sys.stderr_old
                        st.session_state['seq'] = gmmhmm.decode(embs, lengths if st.session_state['use_lengths'] else None)[1]
                        st.session_state['start'] = gmmhmm.startprob_
                        st.session_state['trans'] = gmmhmm.transmat_
                    else:
                        lengths = np.cumsum([0] + (lengths if st.session_state['use_lengths'] else [sum(lengths)]))
                        emis = np.stack([np.exp(gmm.score_samples(embs)) for gmm in gmms])
                        for i in range(max_iter):
                            start0 = np.zeros_like(start)
                            trans0 = np.zeros_like(trans)
                            emis0 = np.zeros_like(emis)
                            p = 0
                            for idx0,idx1 in zip(lengths,lengths[1:]):
                                k = idx1-idx0
                                A = np.zeros((n_topics, k))
                                B = np.zeros((n_topics, k))
                                H = np.zeros((n_topics, n_topics, k))                            

                                for j in range(k):
                                    if j == 0:
                                        A[:,0] = start * emis[:,j+idx0]
                                    else:
                                        A[:,j] = (A[:,j-1] @ trans) * emis[:,j+idx0]
                                    q = A[:,j].max()
                                    A[:,j] /= q
                                    p += np.log(q)
                                p += np.log(A[:,-1].sum())

                                B[:,-1] = 1
                                for j in range(k-2,-1,-1):
                                    B[:,j] = trans @ (B[:,j+1] * emis[:,j+idx0])
                                    B[:,j] /= B[:,j].max()

                                G = A * B;
                                G /= G.sum(axis=0)
                                for j in range(k-1):
                                    H[:,:,j] = (A[:,j:j+1] @ B[:,j+1:j+2].transpose()) * trans * emis[:,idx0+j+1:idx0+j+2].transpose()
                                    H[:,:,j] /= H[:,:,j].sum()

                                start0 += G[:,0]
                                trans0 += H.sum(axis=2)
                                emis0[:,idx0:idx1] = G
                                
                            metrics = {'Перплексия': np.exp(-p/lengths[-1])}
                            st.session_state['progress_msg'] = f'Подгонка кластеров: {i+1}/{max_iter}  \n' + \
                                ''.join(f'{k}: {v:.4f}  \n' for k,v in metrics.items())
                            tm_pbar5.progress((i+1)/max_iter, text=st.session_state['progress_msg'])
                            for t in range(n_topics):
                                gmmfit(gmms[t], embs, emis0[t])

                            start = start0 / start0.sum()
                            trans = trans0 / trans0.sum(axis=1)[...,np.newaxis]
                            emis = np.stack([np.exp(gmm.score_samples(embs)) for gmm in gmms])
                        
                        seq = []
                        for idx0,idx1 in zip(lengths,lengths[1:]):
                            k = idx1-idx0
                            A = start * emis[:,idx0]
                            M = np.zeros((n_topics, k-1), dtype=int)
                            for j in range(1,k):
                                A = A[...,np.newaxis] * trans
                                M[:,j-1] = np.argmax(A, axis=0)
                                A = A.max(axis=0) * emis[:,idx0+j]
                                A /= A.max()
                            res = [int(np.argmax(A))]
                            for j in range(k-2,-1,-1):
                                res = [M[res[0],j]] + res
                            seq += res
                        
                        st.session_state['seq'] = np.array(seq)
                        st.session_state['start'] = start
                        st.session_state['trans'] = trans
                    
                    st.session_state['ready'] = 'SentHMM'
                    fill_result()
                    
                    if use_lib:
                        metrics = {'Полная перплексия': -float(m[-1][1]) if m else 1.0}               
                    metrics = metrics | {'Перплексия': perplexity_hmm(st.session_state['start'], st.session_state['trans'], \
                                                                      st.session_state['phi'].transpose())}
                    metrics = metrics | calc_metrics(st.session_state['counts'].transpose(), st.session_state['phi'], None)
                    st.session_state['progress_msg'] = f'Извлечение тем: {m[-1][0] if use_lib and m else max_iter}/{max_iter}  \n' + \
                            ''.join(f'{k}: {v:.4f}  \n' for k,v in metrics.items())
                    tm_pbar5.progress(1.0, text=st.session_state['progress_msg'])
                
                else:
                    theta = np.zeros((n_topics, len(topics)))
                    for i in range(len(topics)):
                        if topics[i] == -1:
                            theta[0,i] = 1
                        else:
                            theta[:,i] = (1-probs[i]) / (n_topics-1)
                            theta[topics[i]+1,i] = probs[i]
                    phi = (theta @ st.session_state['counts']).transpose()
                    phi = phi / phi.sum(axis=0)
                    
                    counts0 = st.session_state['counts']
                    topics = np.array(st.session_state['topic_model'].topics_)
                    counts = np.vstack([counts0[topics == i-1].sum(axis=0) for i in range(n_topics)])
                    ratio = counts.sum() / counts.sum(axis=0) 
                    scores = (counts/counts.sum(axis=1)[...,np.newaxis]) * np.log(1 + ratio/phi.shape[1])
                    topic_repr = {i-1: [(st.session_state['words'][j], scores[i,j]) for j in np.argsort(scores[i])[::-1][:10]] \
                                  for i in range(n_topics)}
                    
                    st.session_state['ctfidf'] = scores.transpose()
                    st.session_state['topic_model'].topic_representations_ = topic_repr
                    st.session_state['phi'] = phi
                    st.session_state['theta'] = theta
                    st.session_state['ready'] = 'BERTopic'
                    fill_result()
                    
                    metrics = calc_metrics(st.session_state['counts'].transpose(), st.session_state['phi'], st.session_state['theta'])
                    st.session_state['progress_msg'] = f'Извлечение тем: 1/1  \n' + \
                        ''.join(f'{k}: {v:.4f}  \n' for k,v in metrics.items())
                    tm_pbar5.progress(1.0, text=st.session_state['progress_msg'])
            
            else:
                tm_pbar5 = st.progress(1.0 if st.session_state['ready'] == ['BERTopic','SentHMM'] else 0.0, \
                    text=st.session_state['progress_msg'] if st.session_state['ready'] in ['BERTopic','SentHMM'] else empty_msg)
                    
                
        with tab6: # Top2Vec
            n_topics = st.slider('Число тем', min_value=2, max_value=100, value=10, key='n_topics6')
            min_cluster_size = st.slider('Размер кластера', min_value=2, max_value=100, value=15, key='min_cluster_size6')
            min_count = st.slider('Минимальное число вхождений слова', min_value=1, max_value=100, value=10, key='min_count6')
            model_names = ['deeppavlov/rubert-base-cased-sentence',
                           'ai-forever/sbert_large_mt_nlu_ru',
                           'ai-forever/sbert_large_nlu_ru',
                           'ai-forever/ru-en-RoSBERTa',
                           'ai-forever/FRIDA',
                           'doc2vec'
                          ];
            model_name = st.selectbox("Выберите модель", model_names, key='model_name6')
            empty_msg = 'Извлечение тем  \nПерплексия  \nКорреляция  \nКогерентность  \nРазнообразие'
            if st.button('Извлечь темы', key='run_tm6') and docs:
                tm_pbar6 = st.progress(0.0, text='Извлечение тем')
                if 'words' not in st.session_state or len(st.session_state['words']) == 0:
                    prepare_vocab(norm_cb, lemma_cb, pr_pbar)
                                
                umap_args = {'n_neighbors': 15, 'n_components': 5, 'metric': 'cosine', 'random_state': random_state}
                hdbscan_args = {'min_cluster_size': min_cluster_size, 'metric': 'euclidean', 'cluster_selection_method': 'eom'}
                def tokenizer(doc):
                    return [w for w in re.findall(re.compile(token_pattern_filt), doc.lower()) if w not in stopwords]
                documents = [d['input'] for d in docs]
                
                if True:
                    sys.stderr_old = sys.stderr
                    sys.stderr = StringIO()
                    t = ReturnableThread(target=Top2VecNew, kwargs={'embedding_model': model_name, 'documents': documents, \
                                                                    'umap_args':umap_args, 'hdbscan_args':hdbscan_args, \
                                                                    'tokenizer':tokenizer,'min_count':min_count})
                    t.start()
                    while t.is_alive():
                        time.sleep(0.1)
                        text = sys.stderr.getvalue()
                        desc = ['Pre-processing documents for training',
                                'Downloading',
                                'Creating joint document/word embedding',
                                'Creating lower dimension embedding of documents',
                                'Finding dense areas of documents',
                                'Finding topics']
                        msg = ['Предобработка документов для обучения',
                               'Загрузка модели',
                               'Создание совместного эмбеддинга документов/словаря', 
                               'Создание эмбеддингов документов низшей размерности',
                               'Поиск областей документов с высокой плотностью',
                               'Извлечение тем']
                        pos = [text.find(d) for d in desc]
                        if max(pos) >= 0:
                            idx = pos.index(max(pos))
                            tm_pbar6.progress((idx+1)/6, f'{msg[idx]} ({idx+1}/6)')
                    t.join()
                    st.session_state['topic_model'] = t.result
                    sys.stderr = sys.stderr_old
                else:
                    st.session_state['topic_model'] = Top2VecNew(embedding_model=model_name, documents=documents, umap_args=umap_args, \
                                                                 hdbscan_args=hdbscan_args, tokenizer=tokenizer,min_count=min_count)
                             
                if n_topics < len(st.session_state['topic_model'].get_topics()[2]):
                    st.session_state['topic_model'].hierarchical_topic_reduction(n_topics)
                    doc_top = st.session_state['topic_model'].doc_top_reduced
                    topic_word_scores = st.session_state['topic_model'].topic_word_scores_reduced
                    topic_words = st.session_state['topic_model'].topic_words_reduced
                else:
                    n_topics = len(st.session_state['topic_model'].get_topics()[2])
                    doc_top = st.session_state['topic_model'].doc_top
                    topic_word_scores = st.session_state['topic_model'].topic_word_scores
                    topic_words = st.session_state['topic_model'].topic_words
                phi = -10*np.ones((st.session_state['counts'].shape[1], n_topics))
                theta = np.zeros((n_topics, st.session_state['counts'].shape[0]))
                
                for i,t in enumerate(doc_top):
                    theta[t,i] = 1
                for i in range(topic_word_scores.shape[0]):
                    for j in range(topic_word_scores.shape[1]):
                        idx = np.where(st.session_state['words'] == topic_words[i][j])[0]
                        if idx:
                            phi[idx[0],i] = topic_word_scores[i,j]
                phi = scipy.special.softmax(phi, axis=0)
                
                umap_args = {'n_neighbors': 15, 'n_components': 2, 'min_dist': 0.0, 'metric': 'cosine'}
                embeddings_2d = UMAP(**umap_args).fit_transform(st.session_state['topic_model']._get_document_vectors(norm=False)) 
                
                st.session_state['topic_model'].embeddings_2d = embeddings_2d
                st.session_state['phi'] = phi
                st.session_state['theta'] = theta 
                st.session_state['ready'] = 'Top2Vec'
                fill_result()
                
                metrics = calc_metrics(st.session_state['counts'].transpose(), st.session_state['phi'], st.session_state['theta'])
                st.session_state['progress_msg'] = f'Извлечение тем: 1/1  \n' + ''.join(f'{k}: {v:.4f}  \n' for k,v in metrics.items())
                tm_pbar6.progress(1.0, text=st.session_state['progress_msg'])
            else:
                tm_pbar6 = st.progress(1.0 if st.session_state['ready'] == 'Top2Vec' else 0.0, \
                    text=st.session_state['progress_msg'] if st.session_state['ready'] == 'Top2Vec' else empty_msg)
    
    with col2:
        
        tab21,tab22,tab23,tab24,tab25,tab26 = \
        st.tabs(['Темы','Хронология слов','Хронология тем','Кластеризация','Слова по темам','Темы по документам'])
                
        with tab21: # Темы
            if len(st.session_state['words']) > 0:
                words = st.session_state['words'].tolist()
                col21a1,col21a2 = st.columns([0.23,0.77])
                with col21a1:
                    if st.session_state['ready']:
                        topic_sizes = 100*st.session_state['theta'].sum(axis=1) / st.session_state['theta'].shape[1]
                        topicnames = [f'Тема {i+1} ({topic_sizes[i]:.2f}%)' \
                                      for i,_ in enumerate(st.session_state['topicnames'])] + ['Перплексия','По порядку']
                        sub = topicnames.index(st.radio('Темы', topicnames))
                    else:
                        topicnames = ['По частоте','По алфавиту']    
                        sub = topicnames.index(st.radio('Слова', topicnames))
                n_topics = len(topicnames) - 2
                with col21a2:
                    par = f'<br>'
                    for i in range(n_topics):
                        par += f"""<t style="color:rgb(255,255,255);"""
                        tn = st.session_state['topicnames'][i]
                        par += f"""background-color:rgb({255*rgb[i][0]},{255*rgb[i][1]},{255*rgb[i][2]});">{tn}<br>"""
                    st.markdown(par, unsafe_allow_html=True)
                
                col21b1,col21b2,col21b3 = st.columns([0.35,0.3,0.35])
                with col21b1:
                    if sub < n_topics:
                        scores = st.session_state['phi' if st.session_state['ready'] != 'BERTopic' else 'ctfidf'][:,sub]
                        word_idxs = np.argsort(-scores)
                    elif sub == n_topics:
                        scores = st.session_state['counts'].sum(axis=0)
                        word_idxs = np.argsort(-scores)
                    else:
                        scores = st.session_state['counts'].sum(axis=0)
                        word_idxs = np.arange(len(scores))
                    if sub < n_topics:
                        options = [f'{words[i]} ({scores[i]:.4f})' for i in word_idxs]
                    else:
                        options = [f'{words[i]} ({scores[i]})' for i in word_idxs]
                    
                    st.selectbox('Слова темы', options)
                with col21b2:
                    mask = st.session_state['counts'].sum(axis=0) >= st.session_state['word_min_count']
                    words_filt = [f'{s}' for s in st.session_state['words'][mask].tolist()]
                    word = st.selectbox('Коллокант', words_filt)
                with col21b3:
                    values = st.session_state['npmi'][words_filt.index(word)].squeeze()
                    cands = zip(words_filt, values)
                    cands = sorted(cands, key = lambda x: -x[1])
                    st.selectbox('Коллокатор', [f'{w} ({v:.4f})' for w,v in cands])
                
                options = st.session_state['words'].tolist()
                filters = st.multiselect(f'Словарь корпуса ({len(options)})', options) 
                if sub < n_topics:
                    scores = st.session_state['theta'][sub,:]
                elif sub == n_topics:
                    scores = st.session_state['perplexity'] if st.session_state['ready'] else -np.arange(len(docs))
                else:
                    scores = -np.arange(len(docs))
                doc_idxs = np.argsort(-scores)
                
                filt_idxs = [words.index(f) for f in filters]
                doc_idxs = [i for i in doc_idxs if st.session_state['counts'][i,filt_idxs].all()]
                if sub <= n_topics and st.session_state['ready']:
                    docnames = [docs[i]['name'] + f' ({scores[i]:.4f})' for i in doc_idxs]
                else:
                    docnames = [docs[i]['name'] for i in doc_idxs]
                
                col21c1,col21c2 = st.columns([0.85,0.15])
                with col21c1:
                    docname = st.selectbox(f'Отфильтрованные документы ({len(docnames)})', docnames)
                    if docname:
                        idx = docnames.index(docname)
                        idx = doc_idxs[idx]
                    else:
                        idx = None
                with col21c2:
                    if st.session_state['ready'] in markup_cases:
                        st.download_button('Скачать разметку', data=st.session_state['markup'], file_name='Разметка.docx')
                    else:
                        st.button('Скачать разметку', disabled=True)
                
                par = f''
                if not st.session_state['ready'] and idx is not None:
                    par += st.session_state['docs'][idx]['text']
                
                if st.session_state['ready'] and st.session_state['ready'] not in ['HMM','SentHMM'] and idx is not None:                     
                    prev,k0 = 0,-1
                    for i,pos in enumerate(st.session_state['docs'][idx]['seq']):
                        tok = st.session_state['docs'][idx]['tokens'][i]
                        add0 = st.session_state['docs'][idx]['text'][prev:tok.start]
                        temp = st.session_state['phi'][pos,:] * st.session_state['theta'][:,idx]
                        k = np.argmax(temp) % len(rgb)
                        add = st.session_state['docs'][idx]['text'][tok.start:tok.stop]
                        if (temp / temp.sum()).max() > 5/n_topics:
                            if k == k0:
                                par += f"""<t style="color:rgb(255,255,255);"""
                                par += f"""background-color:rgb({255*rgb[k][0]},{255*rgb[k][1]},{255*rgb[k][2]});">{add0+add}</t>"""
                            else:
                                if k0 >= 0 and not any(c.isalpha() for c in add0):
                                    par += f"""<t style="color:rgb(255,255,255);"""
                                    par += f"""background-color:rgb({255*rgb[k0][0]},{255*rgb[k0][1]},{255*rgb[k0][2]});">{add0}</t>"""
                                else:
                                    par += f"""<t>{add0}</t>"""
                                par += f"""<t style="color:rgb(255,255,255);"""
                                par += f"""background-color:rgb({255*rgb[k][0]},{255*rgb[k][1]},{255*rgb[k][2]});">{add}</t>"""
                            k0 = k
                        else:
                            par += f"""<t>{add0+add}</t>"""
                            k0 = -1
                        prev = tok.stop
                    add = st.session_state['docs'][idx]['text'][prev:]
                    par += f"""<t>{add}</t>""" 
                elif st.session_state['ready'] and idx is not None:
                    prev = 0
                    k = 0 if len(st.session_state['labels'][idx]) == 0 else st.session_state['labels'][idx][0]
                    for i,sent in enumerate(st.session_state['docs'][idx]['sents']):
                        add = st.session_state['docs'][idx]['text'][prev:sent.start].replace('.','\.')
                        par += f"""<t style="color:rgb(255,255,255);"""
                        par += f"""background-color:rgb({255*rgb[k][0]},{255*rgb[k][1]},{255*rgb[k][2]});">{add}</t>"""
                        if any(st.session_state['docs'][idx]['parts'] == i):
                            k = mode(st.session_state['labels'][idx][st.session_state['docs'][idx]['parts'] == i]).mode[0] % len(rgb)
                        add = st.session_state['docs'][idx]['text'][sent.start:sent.stop].replace('.','\.')
                        par += f"""<t style="color:rgb(255,255,255);"""
                        par += f"""background-color:rgb({255*rgb[k][0]},{255*rgb[k][1]},{255*rgb[k][2]});">{add}</t>"""
                        prev = sent.stop
                    add = st.session_state['docs'][idx]['text'][prev:]
                    par += f"""<t style="color:rgb(255,255,255);"""
                    par += f"""background-color:rgb({255*rgb[k][0]},{255*rgb[k][1]},{255*rgb[k][2]});">{add}</t>"""
                
                st.markdown(par, unsafe_allow_html=True)
                
        with tab22: # Хронология слов
            if len(st.session_state['words']) > 0:
                dates = [d['date'] for d in st.session_state['docs']]
                dates = [datetime.date(int(d[6:]),int(d[3:5]),int(d[0:2])) for d in dates]
                span = [d.toordinal() for d in dates]
                x = np.arange(span[0], span[-1]+1)
                xt = [datetime.date.fromordinal(k) for k in x]
                options = [f'{w} ({c})' for w,c in zip(st.session_state['words'].tolist(), st.session_state['counts'].sum(axis=0))]
                filters = st.multiselect(f'Словарь корпуса ({len(options)})', options, key='multiselect22') 
                filt_idxs = [options.index(f) for f in filters]
                fig,ax = plt.subplots()
                for idx in filt_idxs:
                    y = np.bincount(span-x[0], weights=st.session_state['counts'][:,idx])
                    y = gaussian_filter1d(y, sigma=10)
                    ax.plot(xt, y)
                ax.legend([st.session_state['words'][idx] for idx in filt_idxs], fontsize=7.5)
                ax.tick_params(axis='both', which='major', labelsize=7.5)
                ax.tick_params(axis='both', which='minor', labelsize=6)
                st.pyplot(fig)
        
        with tab23: # Хронология тем
            if st.session_state['ready']:
                n_topics = st.session_state['phi'].shape[1]
                data = np.zeros(st.session_state['phi'].shape)
                fig, axs = plt.subplots((n_topics + 1)//2, 2, figsize=(13, (n_topics+1)//2*4))
                counts = st.session_state['counts'].sum(axis=1)
                top = (counts * st.session_state['theta']).max()
                for i in range(n_topics):
                    val = counts * st.session_state['theta'][i]
                    tn = st.session_state['topicnames'][i]
                    y = np.bincount(span-x[0], weights=val)
                    y = gaussian_filter1d(y, sigma=10)
                    if n_topics > 2:
                        axs[i//2, i%2].tick_params(axis='both', which='major', labelsize=10)
                        axs[i//2, i%2].tick_params(axis='both', which='minor', labelsize=8)
                        #axs[i//2, i%2].set_ylim(0, top)
                        axs[i//2, i%2].set_title(f'Тема {i}: {tn}')
                        axs[i//2, i%2].plot(xt, y, color=rgb[i])
                    else:
                        axs[i].tick_params(axis='both', which='major', labelsize=10)
                        axs[i].tick_params(axis='both', which='minor', labelsize=8)
                        #axs[i].set_ylim(0, top)
                        axs[i].set_title(f'Тема {i}: {tn}')
                        axs[i].plot(xt, y, color=rgb[i]) 
                st.pyplot(fig)
                
        with tab24: # Кластеризация
            if st.session_state['ready'] in ['LDA','ARTM','HMM']:
                topicnames = [tn[:[a for a in re.finditer(',',tn)][2].span()[0]] for tn in st.session_state['topicnames']]
                color_discrete_map={f'{topicnames.index(t)+1}: {t}':rgb2hex(np.array(r)) for t,r in zip(topicnames, rgb)}
                    
                df = pd.DataFrame(np.hstack([st.session_state['theta_emb'],np.argmax(st.session_state['theta'], axis=0)[...,np.newaxis]]),\
                                  columns=['x','y','z'])
                df['text'] = [d['text'] for d in docs]
                df.sort_values(by='z', inplace=True)
                df['z'] = df['z'].apply(lambda x: f'{int(x)+1}: {topicnames[int(x)]}')
                df['text'] = df['text'].apply(lambda x: x[:100])
                
                fig = px.scatter(
                    df,
                    x='x',
                    y='y',
                    hover_name='text',
                    color='z',
                    height=640,
                    #color_discrete_map=color_discrete_map,
                )
                fig.update_traces(hovertemplate="%{hovertext}<extra></extra>")

                # Display the chart in Streamlit
                st.plotly_chart(fig)
            
            if st.session_state['ready'] == 'CTM':
                ctm_pd = vis.prepare(**st.session_state['lda_vis_data'], n_jobs=1)
                vis.display(ctm_pd)
                html_string = vis.prepared_data_to_html(ctm_pd)
                
                # with open('vis.html','w') as f:
                #     f.write(html_string)
                
                components.v1.html(html_string, width=800, height=800, scrolling=True)
            
            if st.session_state['ready'] == 'BERTopic':
                fig = visualize_documents(st.session_state['topic_model'], [d['text'] for d in docs],
                                          topics = np.array(st.session_state['topic_model'].topics_) + 1,
                                          embeddings=st.session_state['embeddings'],
                                          reduced_embeddings=st.session_state['reduced_embeddings'])
                st.plotly_chart(fig)
                #fig.write_html('topics.html')
            
            if st.session_state['ready'] == 'SentHMM':
                #topics = [0]
                #for d,l in zip(docs,st.session_state['labels']):                    
                #    for i in range(len(d['sents'])):
                #        topics.append(mode(l[d['parts'] == i]).mode[0] if np.any(d['parts'] == i) else topics[-1])      
                #topics = topics[1:]
                topics = st.session_state['seq']
                
                fig = visualize_documents(st.session_state['topic_model'], sum(([s.text for s in d['sents']] for d in docs), []),
                                          topics = topics,
                                          embeddings=st.session_state['embeddings_sent'],
                                          reduced_embeddings=st.session_state['reduced_embeddings_sent'])
                st.plotly_chart(fig)
                #fig.write_html('topics.html')
            
            if st.session_state['ready'] == 'Top2Vec':
                if st.session_state['topic_model'].doc_top_reduced is not None:
                    df = pd.DataFrame(np.hstack([st.session_state['topic_model'].embeddings_2d, \
                                                 st.session_state['topic_model'].doc_top_reduced[...,np.newaxis]]), \
                                 columns=['x','y','z'])
                else:
                    df = pd.DataFrame(np.hstack([st.session_state['topic_model'].embeddings_2d, \
                                                 st.session_state['topic_model'].doc_top[...,np.newaxis]]), \
                                 columns=['x','y','z'])
                    
                topicnames = [tn[:[a for a in re.finditer(',',tn)][2].span()[0]] for tn in st.session_state['topicnames']]
                color_discrete_map={f'{topicnames.index(t)+1}: {t}':rgb2hex(np.array(r)) for t,r in zip(topicnames, rgb)}
                    
                df['text'] = [d['text'] for d in docs]
                df.sort_values(by='z', inplace=True)
                df['z'] = df['z'].apply(lambda x: f'{int(x)+1}: {topicnames[int(x)]}')
                df['text'] = df['text'].apply(lambda x: x[:100])

                fig = px.scatter(
                    df,
                    x='x',
                    y='y',
                    hover_name='text',
                    color='z',
                    height=640,
                    #color_discrete_map=color_discrete_map,
                )
                fig.update_traces(hovertemplate="%{hovertext}<extra></extra>")
                
                st.plotly_chart(fig)
                
        with tab25: # Слова по темам
            if st.session_state['ready'] and 'phi_df' in st.session_state:
                st.dataframe(st.session_state['phi_df'].style.format(precision=4), height=1024)
                
        with tab26: # Темы по документам
            if st.session_state['ready'] and 'theta_df' in st.session_state:
                st.dataframe(st.session_state['theta_df'].style.format(precision=4), height=1024)
        
if __name__ == '__main__':
    try:
        main()
    except SystemExit:
        pass