from tqdm import tqdm
from data import Vocab, wavtoolm2midi, VocabError
import torch.multiprocessing as mp
import sys
import json
import psycopg

# TODO make this into a more flexible tool with Fire
# TODO no reset / select dataset should be options


def reset_lengths(conn):
    with conn.transaction(), conn.cursor() as cur:
        cur.execute("ALTER TABLE extracted_clips DROP COLUMN IF EXISTS symbolic_length")
        cur.execute("ALTER TABLE extracted_clips ADD COLUMN symbolic_length smallint")


def update_lengths_worker(args):
    vocab, dbconnstr, rank, size, limit = args
    with psycopg.connect(dbconnstr) as conn:
        conn.autocommit = True
        whereclause = """
            WHERE c.id %% %s = %s
        """
        with conn.transaction():
            with conn.cursor() as cur:
                cur.execute(
                    "SELECT count(1) FROM extracted_clips c JOIN files f ON c.file_id = f.id "
                    + whereclause,
                    (size, rank),
                )
                total = cur.fetchone()[0]
                cur.execute("SET cursor_tuple_fraction = 0.0001")  # dank
            with conn.cursor(
                name=f"update_lengths_{rank}"
            ) as cur, conn.cursor() as updcur:
                count = 0
                success_count = 0
                cur.itersize = 100

                cur.execute(
                    "SELECT c.id, cn.notes FROM extracted_clips c JOIN extracted_clip_notes cn ON c.id = cn.extracted_clip_id JOIN files f ON c.file_id = f.id "
                    + whereclause
                    + "ORDER BY c.id ASC",
                    (size, rank),
                )

                updlst = []

                def flush():
                    nonlocal updlst
                    updcur.executemany(
                        "UPDATE extracted_clips SET symbolic_length = %s WHERE id = %s",
                        updlst,
                    )
                    updlst = []

                for id, notes in tqdm(cur, total=total, disable=rank != 0):
                    count += 1
                    if rank == 0 and count % 10000 == 0:
                        tqdm.write(f"success rate: {success_count / count}")
                    try:
                        res = vocab.fast_estimate_length(wavtoolm2midi(notes))
                    except VocabError as e:
                        continue
                    if res > 30000 or res <= 5:
                        continue
                    updlst.append((res, id))
                    success_count += 1
                    if len(updlst) >= 1000:
                        flush()
                    if limit is not None and count >= limit:
                        break
                if len(updlst) > 0:
                    flush()


def update_lengths(vocab, parallelism, dbconnstr, limit=None):
    with psycopg.connect(dbconnstr) as conn:
        conn.autocommit = True
        reset_lengths(conn)
        with mp.Pool(parallelism) as pool:
            pool.map(
                update_lengths_worker,
                [(vocab, dbconnstr, i, parallelism, limit) for i in range(parallelism)],
            )
        with conn.cursor() as cur:
            cur.execute("ANALYZE extracted_clips")


if __name__ == "__main__":
    with open(sys.argv[1], "r") as f:
        config = json.load(f)
    vocab = Vocab.from_config(config)
    update_lengths(
        vocab,
        4,
        config["dataset_class"]["conn_str"],
        limit=int(sys.argv[2]) if len(sys.argv) > 2 else None,
    )