import matplotlib.pyplot as mp
import pandas as pd
import seaborn as sb
from settings import settings

from clickhouse_tools import ClickHouseClient
from postgres_tools import PostgresClient


class DataPreparer:
    def __init__(self, clickhouse_host='localhost', postgres_host='localhost', postgres_db='your_db',
                 postgres_user='your_user', postgres_password='your_password'):
        self.clickhouse_client = ClickHouseClient("localhost", 8123, database=settings.DATABASE, username=settings.CLICKHOUSE_USER, password=settings.CLICKHOUSE_PASSWORD)
        self.postgres_client = PostgresClient(
        dbname=settings.DATABASE,
        user=settings.POSTGRES_USER,
        password=settings.POSTGRES_PASSWORD,
        host="localhost",
        port="5432"
    )

    def prepare_ml_dataset(self):
        """
        Подготавливает набор данных для машинного обучения, объединяя данные из ClickHouse и PostgreSQL.
        :return: DataFrame с подготовленными данными.
        """
        # clickhouse_data = self.clickhouse_client.get_data()
        postgres_data = self.postgres_client.get_experiments()
        result_data = self.postgres_client.get_data()

        # Объединение данных по file_id
        ml_dataset = pd.merge(postgres_data, result_data, on='file_id')

        self.postgres_client.close()

        return ml_dataset


data_preparer = DataPreparer()

# Подготовка набора данных для машинного обучения
ml_dataset = data_preparer.prepare_ml_dataset()
ml_dataset = ml_dataset.drop('file_id', axis=1)

ml_dataset.to_csv('burner_data_pg_2.csv', index=False)

# Находим колонки с одним уникальным значением
cols_to_drop = ml_dataset.columns[ml_dataset.nunique() == 1]

# Удаляем эти колонки
ml_dataset = ml_dataset.drop(columns=cols_to_drop)

fig, ax = mp.subplots(figsize=(40, 40))
dataplot = sb.heatmap(ml_dataset.corr(), cmap="YlGnBu", annot=True)

# displaying heatmap
mp.show()