import pymongo
import numpy as np
import time
import pickle
from nltk.corpus import stopwords
import pymorphy3
from sklearn.feature_extraction.text import HashingVectorizer, TfidfTransformer
from sklearn.metrics.pairwise import cosine_similarity
import re

from db import *
from models import *

# Initialize vectorizers and load pre-fitted TfidfTransformer
hashing_vectorizer = HashingVectorizer(
	n_features=2**16,
	alternate_sign=False,
	norm=None,
	binary=False,
	stop_words=None,  # Already removed stop words
	lowercase=False  # Already preprocessed
)

# Load the pre-fitted TfidfTransformer
with open('tfidf_transformer.pkl', 'rb') as f:
	tfidf_transformer = pickle.load(f)

# Initialize Russian stop words and lemmatizer
russian_stop_words = set(stopwords.words('russian'))
morph = pymorphy3.MorphAnalyzer()

# Similarity threshold
SIMILARITY_THRESHOLD = 0.5  # Adjust based on requirements


# Serialization functions
def deserialize_centroid(blob):
	return pickle.loads(blob)


def serialize_centroid(centroid):
	return pickle.dumps(centroid)


def preprocess(clean_text):
	clean_text = clean_text.lower()
	clean_text = re.sub(r"http\S+", "", clean_text)  # Removing URLs

	tokens = clean_text.split()
	lemmatized_tokens = [morph.parse(token)[0].normal_form for token in tokens]
	filtered_tokens = [token for token in lemmatized_tokens if token not in russian_stop_words]
	return ' '.join(filtered_tokens)


def find_best_cluster(tfidf_vector, session, theme):
	statement = select(ThemeStory).where(ThemeStory.theme_id == theme.id)
	clusters = session.exec(statement).all()
	if not clusters:
		return None, 0.0

	centroids = []
	cluster_ids = []
	for cluster in clusters:
		centroid = deserialize_centroid(cluster.centroid)
		centroids.append(centroid)
		cluster_ids.append(cluster.id)

	centroids_matrix = np.array(centroids)
	similarities = cosine_similarity(tfidf_vector, centroids_matrix)
	best_idx = np.argmax(similarities)
	best_similarity = similarities[0, best_idx]

	if best_similarity >= SIMILARITY_THRESHOLD:
		return cluster_ids[best_idx], best_similarity
	else:
		return None, best_similarity


def update_centroid(cluster, new_vector):
	existing_centroid = deserialize_centroid(cluster.centroid)

	# Update centroid by averaging with new vector
	updated_centroid = (existing_centroid + new_vector.toarray()[0]) / 2
	cluster.centroid = serialize_centroid(updated_centroid)


def process_dublicates(theme, session, mdb): #TODO не обгонять антидубль
	statement = select(ThemePost).where(
		ThemePost.theme_id == theme.id, ThemePost.story_id == None, ThemePost.minhash != None
	).order_by(asc(ThemePost.pub_time)).limit(10000)
	theme_posts = session.exec(statement).all()

	for post in theme_posts:
		posts_db = mdb.eye[f"{post.soc_net}_posts"]
		post_meta = posts_db.find_one({"_id": ObjectId(post.post_oid)})
		if not post_meta:
			print(post)
			continue
		post_text = post_meta["text"]

		# Preprocess the content
		processed_content = preprocess(post_text)

		# Vectorize the content
		hashed_vector = hashing_vectorizer.transform([processed_content])
		tfidf_vector = tfidf_transformer.transform(hashed_vector)

		# Find the best cluster
		cluster_id, similarity = find_best_cluster(tfidf_vector, session, theme)

		if cluster_id is not None:
			# Assign to existing cluster
			statement = select(ThemeStory).where(ThemeStory.id == cluster_id)
			cluster = session.scalar(statement)
			update_centroid(cluster, tfidf_vector)
			session.commit()

			assigned_cluster_id = cluster.id

			print(f"Assigned to existing cluster {cluster.id} with similarity {similarity:.2f}")
		else:
			# Create a new cluster
			new_centroid = tfidf_vector.toarray()[0]
			serialized_centroid = serialize_centroid(new_centroid)
			new_cluster = ThemeStory(theme_id=theme.id, title=post_text[:100], centroid=serialized_centroid)
			session.add(new_cluster)
			session.commit()
			assigned_cluster_id = new_cluster.id
			print(f"Created new cluster {new_cluster.id}")

		post.story_id = assigned_cluster_id
		session.execute(update(ThemePost).values({"story_id": assigned_cluster_id}).where(ThemePost.dublicate_of == post.id))
		session.commit()
	print(".. commit fin")


def main():
	# try:
		session = get_session()
		mdb = pymongo.MongoClient("mongodb://eye_writer:5Mo4^hA3c8!saB@192.168.88.3:24317/eye")
		while True:
			session.commit()
			statement = select(Theme).where(Theme.id == 159).order_by(asc(Theme.last_refresh)).limit(1)
			theme = session.exec(statement).one_or_none()
			if not theme:
				time.sleep(1)
			else:
				print(f"got new theme {theme=}")
				process_dublicates(theme, session, mdb)

# except Exception as e:
# 	print("ERRRRRRRRRRRRRR", e)
	# return main()


if __name__ == '__main__':
	main()

make a cache for centroids which will change after theme iteration