import fix_ulimit
from train_target import EMBED

fix_ulimit.fix_ulimit()

import queue
import threading
from dataclasses import dataclass
from typing import List
import torch
from tqdm import tqdm
import dataclasses
import json
import numpy as np
from pgvector.psycopg import register_vector

import typeguard
from model import ModelArgs, MusicalPositionEmbedTransformer
from data import Vocab
from data_gen import DataGenerator, Dataset, TrivialAugmenter
import dataset_classes
import psycopg
import fire
import torch.distributed as dist
from classifier_model import SimpleClassifier


def embed_lanes(config_path: str, model_path: str, batch_size: int):
    dist.init_process_group("nccl")
    rank = dist.get_rank()
    world_size = dist.get_world_size()
    device_id = rank % torch.cuda.device_count()
    device_str = f"cuda:{device_id}"
    torch.cuda.set_device(device_id)

    logging.info(f"start rank {rank+1} of {world_size}")

    with open(config_path, "r") as f:
        config = json.load(f)

    vocab = Vocab.from_config(config)

    conn_str = config["dataset_class"]["conn_str"]

    seq_len = int(config["seq_len"])
    seq_len_min = int(config["seq_len_min"])
    seq_len_max = int(config["seq_len_max"])

    model = MusicalPositionEmbedTransformer(
        vocab,
        dataclasses.replace(
            ModelArgs.from_config(vocab, config),
            max_batch_size=batch_size,
            max_seq_len=seq_len,
            cache=False,
        ),
    )
    model.load_state_dict(
        torch.load(
            model_path,
            map_location=torch.device("cpu"),
        )
    )

    model = model.to(device_str)

    dataset = Dataset.dynamic_from_config(config, "dataset_class")

    data_gen = DataGenerator(
        dataset,
        vocab,
        split=None,
        ranksize=(rank, world_size),
        batch_size=batch_size,
        seq_len=seq_len,
        seq_len_min=seq_len_min,
        seq_len_max=seq_len_max,
        train_target=EMBED,
        parallelism=8,
        to_device=device_str,
        text_tokenize=None,
        example_continuous=True,
        pack_batch=False,
        augmenter=TrivialAugmenter(),
    )

    q_batches = queue.Queue(maxsize=10)
    q_eval_out = queue.Queue(maxsize=100)

    def thd_batchify_entry():
        for batch in tqdm(
            data_gen.generate(order="random"),
            total=int(data_gen.num_examples() / (batch_size * world_size)),
            disable=(rank != 0),
        ):
            q_batches.put(batch)
        q_batches.put(None)

    thd_batchify = threading.Thread(target=thd_batchify_entry, daemon=True)
    thd_batchify.start()

    def thd_eval_entry():
        model.eval()
        with torch.no_grad():
            while True:
                batch = q_batches.get()
                if batch is None:
                    break
                ys, tok_counts = model.forward(
                    batch.tgt, start_pos=0, return_embeddings=True
                )
                q_eval_out.put((batch.ids, ys, tok_counts))
            q_eval_out.put(None)

    thd_eval = threading.Thread(target=thd_eval_entry, daemon=True)
    thd_eval.start()

    @dataclass
    class EmbedRec:
        id: int
        y: np.ndarray
        avg_count: int

    embed_avgs: dict[int, EmbedRec] = dict()

    with psycopg.connect(conn_str) as conn:
        register_vector(conn)
        conn.autocommit = True
        finished_embeds: List[EmbedRec] = []
        with conn.cursor() as cur:

            def flush(to_empty):
                if len(finished_embeds) > (0 if to_empty else 100):
                    with conn.transaction():
                        cur.executemany(
                            """
                            WITH embedding_ids AS (
                                INSERT INTO embeddings (embedding)
                                VALUES (%s)
                                RETURNING %s as lane_id, id as embedding_id
                            ) INSERT INTO lane_embeddings
                                (lane_id, embedding_id)
                                SELECT lane_id, embedding_id FROM embedding_ids
                            """,
                            ((e.y, e.id) for e in finished_embeds),
                        )
                    finished_embeds.clear()

            while True:
                x = q_eval_out.get()
                if x is None:
                    break
                ids, ys, tok_counts = x
                this_batch_embeds = [
                    EmbedRec(id, y * num_toks, num_toks)
                    for id, y, num_toks in zip(
                        ids,
                        ys.cpu().detach().numpy(),
                        tok_counts.cpu().detach().numpy(),
                    )
                    if num_toks > seq_len_min - 1
                ]
                this_batch_ids = set(e.id for e in this_batch_embeds)
                for e in this_batch_embeds:
                    old = embed_avgs.get(e.id)
                    if old is None:
                        embed_avgs[e.id] = e
                    else:
                        old.y += e.y
                        old.avg_count += e.avg_count
                for id in embed_avgs.keys() - this_batch_ids:
                    e = embed_avgs[id]
                    del embed_avgs[id]
                    e.y /= e.avg_count
                    e.avg_count = 1
                    finished_embeds.append(e)
                flush(False)
            flush(True)

    thd_batchify.join()
    thd_eval.join()


def scratch():
    # gather some lane-embeddings that are either drums or not drums
    with psycopg.connect("host=localhost dbname=composer_new_dataset_v3") as conn:
        register_vector(conn)
        cur = conn.cursor()
        cur.execute(
            """
            select e.embedding, exists(
                    select from lane_tags lt
                    join tag_values tv on lt.tag_value_id = tv.id
                    where lt.lane_id = le.lane_id
                    and tv.value in ('Flagged Drums', 'Analyzed Drums', 'ML Analyzed Drums')
            ) as is_drum
            from embeddings e
            join lane_embeddings le on e.id = le.embedding_id
            join lanes l on l.id = le.lane_id
            join instruments i on i.id = l.instrument_id
            join files f on f.id = i.file_id
            and exists(
                select from lane_tags lt
                join tag_values tv on lt.tag_value_id = tv.id
                where lt.lane_id = le.lane_id
                and tv.value in ('Flagged Drums', 'Analyzed Drums', 'ML Analyzed Drums', 'ML Analyzed Non-Drums')
            )
            order by f.rand_order
            limit 100000
            """
        )
        xs = cur.fetchall()
        drums = [x for x, is_drum in xs if is_drum]
        not_drums = [x for x, is_drum in xs if not is_drum]
        # not_drums = random.sample(not_drums, len(drums))

        # train a SVM classifier
        from sklearn.svm import SVC
        from sklearn.model_selection import train_test_split
        from sklearn.metrics import classification_report

        X = np.array(drums + not_drums)
        y = np.array([1] * len(drums) + [0] * len(not_drums))
        # y = np.random.permutation(y)  # confuse
        model = SVC(class_weight="balanced")
        X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2)

        print("baseline:")
        print(classification_report(y_test, [0] * len(y_test)))

        # model.fit(X_train, y_train)
        # y_pred = model.predict(X_test)
        # print("svm:")
        # print(classification_report(y_test, y_pred))

        model2 = SimpleClassifier()
        optim = torch.optim.SGD(model2.parameters(), lr=0.01)
        loss = torch.nn.BCELoss()
        for i in range(1000):
            drum_idxs = np.nonzero(y_train == 1)[0]
            not_drum_idxs = np.nonzero(y_train == 0)[0]
            some_drums = random.sample(drum_idxs.tolist(), 16)
            some_not_drums = random.sample(not_drum_idxs.tolist(), 16)
            X_epoch = torch.tensor(
                np.concatenate([X_train[some_drums], X_train[some_not_drums]])
            )
            y_epoch = torch.tensor([1] * 16 + [0] * 16)
            y_pred = model2(X_epoch).squeeze()
            l = loss(y_pred, y_epoch.float())
            optim.zero_grad()
            l.backward()
            optim.step()

        with torch.no_grad():
            y_pred = model2.forward(torch.tensor(X_test)).squeeze().numpy() > 0.90
        print("logistic:")
        print(classification_report(y_test, y_pred))


def parse_bool(s):
    if s.lower() in ("true", "1"):
        return True
    elif s.lower() in ("false", "0"):
        return False
    else:
        raise ValueError(f"Cannot convert to bool: {s}")


class ClassifierTool:
    @typeguard.typechecked
    def embed_lanes(
        self,
        config_path: str,
        model_path: str,
        batch_size: int = 16,
    ):
        """
        Generate and write to the database embedding vectors for all lanes lacking them

        Model and vocab config will be read from config_path
        Model will be loaded from model_path, which should be stripped (see loader.py)
        Database connection string will be read from config["dataset_class"]["conn_str"]

        Should be invoked with torchrun --standalone --nproc-per-node gpu
        """
        embed_lanes(config_path, model_path, batch_size)

    def scratch(self):
        scratch()

    # @fire.decorators.SetParseFn(str, "class_tags", "nonclass_tags")
    # @typeguard.typechecked
    # def collect(
    #     self,
    #     class_tags: str,
    #     nonclass_tags: str,
    #     class_limit: Optional[int] = None,
    #     nonclass_multiplier: int = 1,
    #     examples_out: str = "data.pt",
    # ):
    #     """
    #     Collects examples for training the classifier.
    #     """
    #     collect_examples(
    #         [TagFilter.from_str(s) for s in class_tags.split(",")],
    #         class_limit,
    #         [TagFilter.from_str(s) for s in nonclass_tags.split(",")],
    #         nonclass_multiplier,
    #         examples_out,
    #     )

    # @typeguard.typechecked
    # def train(
    #     self,
    #     examples: str = "data.pt",
    #     batch_size: int = 16,
    #     train_epochs: int = 15,
    #     classifier_out: str = "drum_classifier.pt",
    #     val_split: float = 0.2,
    # ):
    #     """
    #     Trains the classifier.
    #     """
    #     train(examples, batch_size, train_epochs, classifier_out, val_split)

    # @fire.decorators.SetParseFn(str, "high_category_tag", "low_category_tag")
    # @fire.decorators.SetParseFn(parse_bool, "lane_mode")
    # @typeguard.typechecked
    # def classify(
    #     self,
    #     classifier: str = "drum_classifier.pt",
    #     batch_size: int = 100,
    #     high_category_tag: Optional[str] = "ML Analyzed Drums",
    #     low_category_tag: Optional[str] = "ML Analyzed Non-Drums",
    #     threshold: float = 0.5,
    #     lane_mode: bool = True,
    #     max_clips_per_lane: int = 8,
    #     limit: Optional[int] = None,
    # ):
    #     """
    #     Classifies lanes.
    #     """
    #     classify(
    #         classifier,
    #         batch_size,
    #         high_category_tag,
    #         low_category_tag,
    #         threshold,
    #         lane_mode,
    #         max_clips_per_lane,
    #         limit,
    #     )


if __name__ == "__main__":
    from util import configure_logging

    configure_logging()

    fire.Fire(ClassifierTool())