from pathlib import Path

import numpy as np
import pandas as pd
import torch

from durbango import lmap, tqdm_nice
from transformers import BartTokenizer


try:
    from finetune import calculate_rouge
except ImportError:
    from .finetune import calculate_rouge


LOGDIR = Path("examples/summarization/bart/dbart/logs/").absolute()
DATA_DIR = Path("examples/summarization/bart/dbart/cnn_dm").absolute()


def rouge_files(src_file: Path, tgt_file: Path):
    src = lmap(str.strip, list(src_file.open().readlines()))
    tgt = lmap(str.strip, list(tgt_file.open().readlines()))
    return calculate_rouge(src, tgt)


def read_gens(exp_name, split="test", n=None):
    if Path(exp_name).exists():
        expdir = exp_name
    else:
        expdir = LOGDIR / exp_name
    assert expdir.exists(), expdir
    paths = list(expdir.glob(f"{split}_generations*.txt"))
    assert paths
    path = paths[0]
    lns = lmap(str.strip, list(path.open().readlines()))
    if n is not None:
        return lns[:n]
    return lns


def load_cpu(p):
    return torch.load(p, map_location="cpu")


class RougeTracker:
    def __init__(self, csv_path="rouge_test_df.csv", logdir=LOGDIR, data_dir=DATA_DIR):
        try:
            self.df = pd.read_csv(csv_path, index_col=0)
        except FileNotFoundError:
            self.df = pd.DataFrame()
        self.logdir = logdir
        test_gt = lmap(str.strip, Path(data_dir / "test.target").open().readlines())
        test_gt = lmap(str.strip, test_gt)
        self.gt = test_gt  # {'test': test_gt}
        self.tokenizer = BartTokenizer.from_pretrained("facebook/bart-large-cnn")
        self.csv_path = csv_path
        self.new_results = pd.DataFrame()

    @property
    def finished_experiments(self):
        return [p.parent.name for p in list(self.logdir.glob("*/test_generations*.txt"))]

    def tok_len(self, strang):
        return len(self.tokenizer.encode(strang))

    def read_all_gens(self,):
        GENS = {}
        for f in self.finished_experiments:
            GENS[f] = read_gens(f)
        return GENS

    def score(self, gens, k, all_stats=False):
        rouge_raw = calculate_rouge(gens, self.gt, all_stats=all_stats)
        lens = np.mean(lmap(self.tok_len, gens))
        return dict(avg_len=lens, exp_name=k, **rouge_raw)

    @property
    def new_experiments(self):
        possible = set(self.finished_experiments).difference(self.df.index)
        to_score = set()
        for p in possible:
            gens = read_gens(p)
            if len(gens) == len(self.gt):
                to_score.add(p)
        return to_score

    def update(self):
        records = []
        to_score = self.new_experiments
        for exp_name in tqdm_nice(to_score, desc="Rouge Update"):
            gens = read_gens(exp_name)
            if len(gens) != len(self.gt):
                continue
            records.append(self.score(gens, exp_name))
        if not records:
            return self.df
        new_df = (
            pd.DataFrame(records).rename(columns=lambda x: x.replace("rouge", "R")).set_index("exp_name").astype(float)
        )
        self.new_results = new_df
        self.df = pd.concat([self.df, new_df]).dsort("R2")
        return self.df