#!/usr/bin/env python
# coding: utf-8

import io
import csv
import multiprocessing as mp
import os
import pickle
import shutil
import sys
import tempfile
import warnings
from collections import Counter, defaultdict
from multiprocessing.pool import ThreadPool
from typing import Any, Dict, List, Optional, Tuple

import numpy as np
import pandas as pd
from conllu import parse, parse_incr

__all__ = ['DepValAnalyzer','analyze','Converter','convert']

punctdeps = ['punkt','punct','pu','PU','PUN','WP','wp']
rootdeps = ['root','ROOT','s','HED','']
stopdeps = ['punct','punkt','_','PUN','PU','pu','wp','WP']

_MP_UNAVAILABLE_WARNED = False


def _can_use_multiprocessing() -> bool:
    # Windows spawn needs a real, importable __main__ file.
    main_mod = sys.modules.get('__main__')
    main_file = getattr(main_mod, '__file__', None) if main_mod is not None else None
    if not main_file:
        return False
    if '<stdin>' in str(main_file):
        return False
    return os.path.isfile(main_file)


def _open_parallel_pool(processes: int):
    global _MP_UNAVAILABLE_WARNED
    if _can_use_multiprocessing():
        return mp.Pool(processes=processes)
    if not _MP_UNAVAILABLE_WARNED:
        warnings.warn(
            "Process-based multiprocessing is unavailable in the current launch context; "
            "falling back to ThreadPool parallel execution.",
            RuntimeWarning,
            stacklevel=2
        )
        _MP_UNAVAILABLE_WARNED = True
    return ThreadPool(processes=processes)


def _resolve_jobs(parallel: bool, n_jobs: Optional[int], n_items: int) -> Tuple[bool, int]:
    if not parallel or n_items < 2:
        return False, 1
    cpu = os.cpu_count() or 1
    if n_jobs is None:
        workers = cpu
    else:
        if n_jobs <= 0:
            raise ValueError(f"n_jobs must be a positive integer, got {n_jobs}.")
        workers = min(n_jobs, cpu)
    if workers <= 1 or n_items < 8:
        return False, 1
    return True, workers


def _recommended_chunksize(n_items: int, workers: int) -> int:
    if workers <= 1:
        return 1
    return max(1, min(256, n_items // max(1, workers * 4)))


def _flatten(values: List[List[Any]]) -> List[Any]:
    return [item for sub in values for item in sub]


def _safe_root_distance(dd_values: List[int]) -> int:
    try:
        return dd_values.index(0) + 1
    except ValueError:
        return 0


def _safe_twig(hd_values: List[int]) -> int:
    if not hd_values:
        return 0
    mode_value = Counter(hd_values).most_common(1)[0][0]
    return mode_value + 1


def _compute_hd(word_id: int, head_map: Dict[int, int], cache: Dict[int, int]) -> int:
    if word_id in cache:
        return cache[word_id]
    current = word_id
    depth = 0
    in_path = set()
    size = max(1, len(head_map))
    while True:
        if current in in_path:
            depth = 0
            break
        in_path.add(current)
        head = head_map.get(current)
        if head is None:
            depth = 0
            break
        if head == 0:
            break
        depth += 1
        if depth > size:
            depth = 0
            break
        if head in cache:
            depth += cache[head]
            break
        current = head
    cache[word_id] = depth
    return depth


def _compute_sentence_dep(payload: Tuple[List[Dict[str, Any]], Tuple[str, ...], Tuple[str, ...], Tuple[str, ...]]) -> Dict[str, List[Any]]:
    sentence, metrics, stop_deprels, root_deprels = payload
    metric_set = set(metrics)
    result: Dict[str, List[Any]] = {metric: [] for metric in metrics}

    valid_words = [
        word for word in sentence
        if isinstance(word.get('id'), int) and word.get('deprel') not in stop_deprels
    ]
    word_index_by_id = {word['id']: word for word in valid_words}
    head_map = {word['id']: int(word.get('head', 0) or 0) for word in valid_words}
    hd_cache: Dict[int, int] = {}
    child_counts = Counter(
        int(word.get('head', 0) or 0)
        for word in valid_words
        if isinstance(word.get('head'), int)
    )

    for word in valid_words:
        wid = word['id']
        head = int(word.get('head', 0) or 0)
        deprel = word.get('deprel')
        upos = word.get('upos')

        is_root = head == 0
        include_non_root = (
            not is_root
            and deprel not in root_deprels
            and upos != 'pu'
        )
        include_root = is_root
        if not (include_non_root or include_root):
            continue

        if 'id' in metric_set:
            result['id'].append(wid)
        if 'form' in metric_set:
            result['form'].append(word.get('form'))
        if 'dpos' in metric_set:
            result['dpos'].append(upos)
        if 'head' in metric_set:
            result['head'].append(head)
        if 'deprel' in metric_set:
            result['deprel'].append(deprel)

        if 'gform' in metric_set:
            if is_root:
                result['gform'].append('ROOT')
            else:
                head_word = word_index_by_id.get(head)
                result['gform'].append(head_word.get('form') if head_word else 'ERROR_ANNOTATION')
        if 'gpos' in metric_set:
            if is_root:
                result['gpos'].append('ROOT')
            else:
                head_word = word_index_by_id.get(head)
                result['gpos'].append(head_word.get('upos') if head_word else 'ERROR_ANNOTATION')

        if 'dd' in metric_set:
            result['dd'].append(0 if is_root else abs(head - wid))
        if 'hd' in metric_set:
            result['hd'].append(0 if is_root else _compute_hd(wid, head_map, hd_cache))
        if 'ddir' in metric_set:
            if is_root:
                result['ddir'].append(0)
            else:
                result['ddir'].append(1 if head < wid else -1)
        if 'v' in metric_set:
            result['v'].append(child_counts.get(wid, 0) if is_root else child_counts.get(wid, 0) + 1)

    return result


def _compute_sentence_pvp(
    payload: Tuple[List[Dict[str, Any]], Optional[str], str, Tuple[str, ...]]
) -> Tuple[Dict[str, int], Dict[str, int]]:
    sentence, input_upos, target, stop_deprels = payload

    valid_words = [
        word for word in sentence
        if isinstance(word.get('id'), int) and word.get('deprel') not in stop_deprels
    ]
    by_id = {word['id']: word for word in valid_words}
    children = defaultdict(list)
    for word in valid_words:
        head = word.get('head', 0)
        if isinstance(head, int):
            children[head].append(word)

    dependents = Counter()
    governors = Counter()
    for word in valid_words:
        if input_upos is not None and word.get('upos') != input_upos:
            continue
        wid = word['id']
        if target == 'deprel':
            for dep in children.get(wid, []):
                dependents[dep.get('deprel')] += 1
            governors[word.get('deprel')] += 1
        elif target == 'pos':
            for dep in children.get(wid, []):
                dependents[dep.get('upos')] += 1
            gov = by_id.get(word.get('head'))
            if gov is not None:
                governors[gov.get('upos')] += 1

    return dict(dependents), dict(governors)


def _calculate_pvp_from_dep_data(
    dep_data: Dict[str, List[List[Any]]],
    input_upos: Optional[str] = None,
    target: str = 'deprel',
    normalize: bool = True
) -> Dict[str, List[Tuple[Any, float]]]:
    dependents = Counter()
    governors = Counter()

    sent_count = len(dep_data.get('id', []))
    for sidx in range(sent_count):
        ids = dep_data['id'][sidx]
        heads = dep_data['head'][sidx]
        dpos = dep_data['dpos'][sidx]
        deprels = dep_data['deprel'][sidx]
        idx_by_id = {wid: i for i, wid in enumerate(ids)}
        children = defaultdict(list)
        for i, head in enumerate(heads):
            children[head].append(i)

        for i, wid in enumerate(ids):
            if input_upos is not None and dpos[i] != input_upos:
                continue

            if target == 'deprel':
                for child_i in children.get(wid, []):
                    dependents[deprels[child_i]] += 1
                governors[deprels[i]] += 1
            elif target == 'pos':
                for child_i in children.get(wid, []):
                    dependents[dpos[child_i]] += 1
                h = heads[i]
                hidx = idx_by_id.get(h)
                if hidx is not None:
                    governors[dpos[hidx]] += 1

    total_deps = sum(dependents.values())
    total_govs = sum(governors.values())
    if normalize:
        govs = sorted(
            ((k, (v / total_govs) if total_govs else 0) for k, v in governors.items()),
            key=lambda x: x[1],
            reverse=True
        )
        deps = sorted(
            ((k, (v / total_deps) if total_deps else 0) for k, v in dependents.items()),
            key=lambda x: x[1],
            reverse=True
        )
    else:
        govs = sorted(governors.items(), key=lambda x: x[1], reverse=True)
        deps = sorted(dependents.items(), key=lambda x: x[1], reverse=True)
    return {'act as a gov': govs, 'act as a dep': deps}


def _estimate_sentence_count(path: str) -> int:
    # CoNLLU sentences are separated by blank lines; this is a fast estimate for scheduling.
    count = 0
    in_sent = False
    with open(path, encoding='utf-8') as f:
        for line in f:
            if line.strip():
                in_sent = True
            elif in_sent:
                count += 1
                in_sent = False
    if in_sent:
        count += 1
    return count


def _select_parallel_mode(
    parallel: bool,
    n_jobs: Optional[int],
    file_stats: List[Dict[str, Any]]
) -> Tuple[str, int]:
    large_treebank_threshold = 5000
    workers = (os.cpu_count() or 1) if n_jobs is None else n_jobs
    if workers <= 0:
        raise ValueError(f"n_jobs must be a positive integer, got {n_jobs}.")
    workers = min(workers, os.cpu_count() or 1)
    if not parallel or workers <= 1:
        return 'serial', 1

    if len(file_stats) == 1:
        # Single-file workloads should use sentence-level parallelism.
        return 'sentence-level', workers

    large_count = sum(1 for x in file_stats if int(x.get('sentence_count', 0)) > large_treebank_threshold)
    if large_count >= 2:
        # When there are multiple large treebanks, prioritize in-file parallelism.
        return 'sentence-level', workers

    n_files = len(file_stats)
    sentence_counts = [max(1, int(x['sentence_count'])) for x in file_stats]
    mean_count = sum(sentence_counts) / max(1, n_files)
    is_uniform = (max(sentence_counts) / max(1, mean_count)) <= 1.5
    if n_files >= 2 * workers and is_uniform:
        return 'file-level', workers
    return 'sentence-level', workers


def _analyze_single_file_worker(task: Dict[str, Any]) -> Tuple[int, str, Dict[str, float]]:
    with open(task['treebank_path'], encoding='utf-8') as treebank:
        text_metrics, dep_metrics, sent_metrics, pvp_metrics, distribution_dict = getDepValFeatures(
            treebank,
            normalize=task['normalize'],
            parallel=False,
            n_jobs=1,
            punct_deprels=task['punct_deprels'],
            root_deprels=task['root_deprels']
        )

    os.makedirs(task['out_dir'], exist_ok=True)
    dep_metrics.round(2).to_csv(os.path.join(task['out_dir'], 'dep_metrics.csv'), index=False)
    sent_metrics.round(2).to_csv(os.path.join(task['out_dir'], 'sent_metrics.csv'), index=False)
    pvp_metrics.round(4).to_csv(os.path.join(task['out_dir'], 'pvp.csv'), index=False)
    for metric_name, metric_df in distribution_dict.items():
        metric_df.to_csv(os.path.join(task['out_dir'], f'{metric_name}_distribution.csv'), index=False, header=False)

    return task['index'], task['file_name'], text_metrics


def _split_batches(tasks: List[Dict[str, Any]], batch_size: int) -> List[List[Dict[str, Any]]]:
    return [tasks[i:i + batch_size] for i in range(0, len(tasks), batch_size)]


def _analyze_file_batch_worker(batch: List[Dict[str, Any]]) -> List[Tuple[int, str, Dict[str, float]]]:
    return [_analyze_single_file_worker(task) for task in batch]


def _analyze_single_file_serial(
    treebank_path: str,
    out_dir: str,
    normalize: bool,
    punct_deprels: Optional[List[str]],
    root_deprels: Optional[List[str]]
) -> Dict[str, float]:
    task = {
        'index': 0,
        'treebank_path': treebank_path,
        'file_name': os.path.basename(treebank_path).split('.')[0],
        'out_dir': out_dir,
        'normalize': normalize,
        'punct_deprels': punct_deprels,
        'root_deprels': root_deprels,
    }
    _, _, text_metrics = _analyze_single_file_worker(task)
    return text_metrics


def _analyze_single_file_sentence_level(
    treebank_path: str,
    out_dir: str,
    normalize: bool,
    workers: int,
    punct_deprels: Optional[List[str]],
    root_deprels: Optional[List[str]]
) -> Dict[str, float]:
    with open(treebank_path, encoding='utf-8') as treebank:
        analyzer = DepValAnalyzer(
            treebank,
            parallel=False,
            n_jobs=1,
            punct_deprels=punct_deprels,
            root_deprels=root_deprels
        )
    sent_offsets = np.zeros(len(analyzer.treebank) + 1, dtype=np.int64)
    if len(analyzer.treebank) > 0:
        sent_offsets[1:] = np.cumsum([len(s) for s in analyzer.treebank], dtype=np.int64)
    ranges = _build_sentence_ranges(sent_offsets, workers)
    part_dir = tempfile.mkdtemp(prefix='depval_parts_')
    tasks = []
    for gid, (st, ed) in enumerate(ranges):
        tasks.append({
            'group_id': gid,
            'start_sent': st,
            'sentences': analyzer.treebank[st:ed],
            'part_dir': part_dir,
            'punct_deprels': list(analyzer.punct_deprels),
            'root_deprels': list(analyzer.root_deprels),
        })
    results = []
    try:
        with _open_parallel_pool(processes=workers) as pool:
            for r in pool.imap_unordered(_sentence_group_to_parts_worker, tasks, chunksize=1):
                results.append(r)
        results = sorted(results, key=lambda x: x['group_id'])
        dep_part_paths = [r['dep_part_path'] for r in results]
        sent_part_paths = [r['sent_part_path'] for r in results]
        _concat_csv_parts(dep_part_paths, os.path.join(out_dir, 'dep_metrics.csv'))
        _concat_csv_parts(sent_part_paths, os.path.join(out_dir, 'sent_metrics.csv'))

        dep_counters = {k: Counter() for k in ['dd', 'hd', 'v', 'deprel', 'pos']}
        sent_counters = {k: Counter() for k in ['sl', 'tw', 'th', 'rd']}
        pvp_dep = Counter()
        pvp_gov = Counter()
        for r in results:
            for k in dep_counters:
                dep_counters[k].update(r['dep_counters'][k])
            for k in sent_counters:
                sent_counters[k].update(r['sent_counters'][k])
            pvp_dep.update(r['pvp_dep'])
            pvp_gov.update(r['pvp_gov'])

        pvp_metrics = _pvp_df_from_counters(pvp_dep, pvp_gov)
        pvp_metrics.round(4).to_csv(os.path.join(out_dir, 'pvp.csv'), index=False)
        distribution_dict = _distribution_from_counters(dep_counters, sent_counters, normalize=normalize)
        for m in distribution_dict:
            distribution_dict[m].to_csv(os.path.join(out_dir, f'{m}_distribution.csv'), index=False, header=False)
        text_metrics = _text_metrics_from_partials([r['text_partial'] for r in results])
    finally:
        shutil.rmtree(part_dir, ignore_errors=True)
    return text_metrics


def _encode_group_results(
    group_results: List[Dict[str, List[Any]]]
) -> Tuple[List[Dict[str, List[Any]]], Dict[str, List[str]]]:
    string_metrics = {'form', 'dpos', 'gform', 'gpos', 'deprel'}
    vocabs: Dict[str, List[str]] = {}
    encoded: List[Dict[str, List[Any]]] = []
    for sent in group_results:
        encoded_sent: Dict[str, List[Any]] = {}
        for metric, values in sent.items():
            if metric in string_metrics:
                vocab = vocabs.setdefault(metric, [])
                table = {v: i for i, v in enumerate(vocab)}
                code_values: List[int] = []
                for val in values:
                    sval = str(val)
                    if sval not in table:
                        table[sval] = len(vocab)
                        vocab.append(sval)
                    code_values.append(table[sval])
                encoded_sent[metric] = code_values
            else:
                encoded_sent[metric] = values
        encoded.append(encoded_sent)
    return encoded, vocabs


def _decode_group_results(
    encoded_group_results: List[Dict[str, List[Any]]],
    vocabs: Dict[str, List[str]]
) -> List[Dict[str, List[Any]]]:
    decoded: List[Dict[str, List[Any]]] = []
    for sent in encoded_group_results:
        decoded_sent: Dict[str, List[Any]] = {}
        for metric, values in sent.items():
            if metric in vocabs:
                vocab = vocabs[metric]
                decoded_sent[metric] = [vocab[int(v)] for v in values]
            else:
                decoded_sent[metric] = values
        decoded.append(decoded_sent)
    return decoded


def _sentence_group_worker_to_file(
    task: Tuple[int, int, List[List[Dict[str, Any]]], Tuple[str, ...], Tuple[str, ...], Tuple[str, ...], str]
) -> Tuple[int, int, str, int, int, int]:
    group_id, start_sent, sentence_chunk, metrics, stop_deprels, root_deprels, temp_dir = task
    results = []
    for sent in sentence_chunk:
        payload = (sent, metrics, stop_deprels, root_deprels)
        results.append(_compute_sentence_dep(payload))
    encoded_results, vocabs = _encode_group_results(results)
    part_path = os.path.join(temp_dir, f'group_{group_id:06d}.pkl')
    with open(part_path, 'wb') as f:
        pickle.dump({'vocabs': vocabs, 'results': encoded_results}, f, protocol=pickle.HIGHEST_PROTOCOL)
    token_count = sum(len(s) for s in sentence_chunk)
    return group_id, start_sent, part_path, os.path.getsize(part_path), len(sentence_chunk), token_count


def _sentence_group_to_parts_worker(task: Dict[str, Any]) -> Dict[str, Any]:
    analyzer = DepValAnalyzer(
        task['sentences'],
        cleaned=True,
        parallel=False,
        n_jobs=1,
        punct_deprels=task['punct_deprels'],
        root_deprels=task['root_deprels']
    )
    dep_data = analyzer.calculate_dep_metrics(parallel=False, n_jobs=1)
    sent_data = analyzer._calculate_sent_metrics_from_dep(dep_data, analyzer.sent_metrics)

    dep_part_path = os.path.join(task['part_dir'], f"dep_{task['group_id']:06d}.csv")
    sent_part_path = os.path.join(task['part_dir'], f"sent_{task['group_id']:06d}.csv")
    dep_columns = list(dep_data.keys())
    with open(dep_part_path, 'w', encoding='utf-8', newline='') as f:
        writer = csv.writer(f)
        writer.writerow(dep_columns)
        sent_count = len(dep_data['id'])
        for si in range(sent_count):
            token_count = len(dep_data['id'][si])
            for ti in range(token_count):
                writer.writerow([dep_data[col][si][ti] for col in dep_columns])

    pd.DataFrame(sent_data).round(2).to_csv(sent_part_path, index=False)

    dep_counters = {
        'dd': Counter(v for v in _flatten(dep_data['dd']) if v > 0),
        'hd': Counter(v for v in _flatten(dep_data['hd']) if v > 0),
        'v': Counter(v for v in _flatten(dep_data['v']) if v > 0),
        'deprel': Counter(_flatten(dep_data['deprel'])),
        'pos': Counter(_flatten(dep_data['dpos'])),
    }
    pvp = _calculate_pvp_from_dep_data(dep_data, input_upos=None, target='deprel', normalize=False)
    sl_values = sent_data['sl']
    den_dep_raw = sum(max(0, int(v) - 1) for v in sl_values)
    text_partial = {
        'n_sent': len(sl_values),
        'sl_sum': float(sum(sl_values)),
        'dep_den_raw': float(den_dep_raw),
        'tdl_sum': float(sum(sent_data['tdl'])),
        'mhd_num': float(sum(mhd * max(0, sl - 1) for mhd, sl in zip(sent_data['mhd'], sl_values))),
        'mv_num': float(sum(mv * sl for mv, sl in zip(sent_data['mv'], sl_values))),
        'vk_sum': float(sum(sent_data['vk'])),
        'tw_sum': float(sum(sent_data['tw'])),
        'th_sum': float(sum(sent_data['th'])),
        'mhdd_sum': float(sum(sent_data['mhdd'])),
        'ndd_sum': float(sum(sent_data['ndd'])),
        'hi_sum': float(sum(sent_data['hi'])),
        'hf_sum': float(sum(sent_data['hf'])),
        'rd_sum': float(sum(sent_data['rd'])),
    }
    sent_counters = {
        'sl': dict(Counter(int(v) for v in sent_data['sl'])),
        'tw': dict(Counter(int(v) for v in sent_data['tw'])),
        'th': dict(Counter(int(v) for v in sent_data['th'])),
        'rd': dict(Counter(int(v) for v in sent_data['rd'])),
    }
    return {
        'group_id': task['group_id'],
        'start_sent': task['start_sent'],
        'dep_part_path': dep_part_path,
        'sent_part_path': sent_part_path,
        'dep_part_size': os.path.getsize(dep_part_path),
        'sent_part_size': os.path.getsize(sent_part_path),
        'sent_count': len(task['sentences']),
        'token_count': int(sum(len(s) for s in task['sentences'])),
        'dep_counters': {k: dict(v) for k, v in dep_counters.items()},
        'sent_counters': sent_counters,
        'text_partial': text_partial,
        'pvp_dep': dict(pvp['act as a dep']),
        'pvp_gov': dict(pvp['act as a gov']),
    }


def _concat_csv_parts(part_paths: List[str], out_path: str) -> None:
    if not part_paths:
        pd.DataFrame().to_csv(out_path, index=False)
        return
    with open(out_path, 'w', encoding='utf-8', newline='') as fout:
        for i, p in enumerate(part_paths):
            with open(p, 'r', encoding='utf-8') as fin:
                if i == 0:
                    fout.write(fin.read())
                else:
                    _ = fin.readline()
                    fout.write(fin.read())


def _text_metrics_from_partials(partials: List[Dict[str, float]]) -> Dict[str, float]:
    if not partials:
        return {k: 0.0 for k in ['mdd','ndd','mhd','mhdd','mtdl','msl','mv','vk','mtw','mth','hi','hf','mrd']}
    n_sent = sum(p['n_sent'] for p in partials)
    if n_sent <= 0:
        return {k: 0.0 for k in ['mdd','ndd','mhd','mhdd','mtdl','msl','mv','vk','mtw','mth','hi','hf','mrd']}
    sl_sum = sum(p['sl_sum'] for p in partials)
    den_dep = max(1.0, sum(p['dep_den_raw'] for p in partials))
    mv_den = max(1.0, sl_sum)
    return {
        'mdd': float(sum(p['tdl_sum'] for p in partials) / den_dep),
        'ndd': float(sum(p['ndd_sum'] for p in partials) / n_sent),
        'mhd': float(sum(p['mhd_num'] for p in partials) / den_dep),
        'mhdd': float(sum(p['mhdd_sum'] for p in partials) / n_sent),
        'mtdl': float(sum(p['tdl_sum'] for p in partials) / n_sent),
        'msl': float(sl_sum / n_sent),
        'mv': float(sum(p['mv_num'] for p in partials) / mv_den),
        'vk': float(sum(p['vk_sum'] for p in partials) / n_sent),
        'mtw': float(sum(p['tw_sum'] for p in partials) / n_sent),
        'mth': float(sum(p['th_sum'] for p in partials) / n_sent),
        'hi': float(sum(p['hi_sum'] for p in partials) / den_dep),
        'hf': float(sum(p['hf_sum'] for p in partials) / den_dep),
        'mrd': float(sum(p['rd_sum'] for p in partials) / n_sent),
    }


def _distribution_from_counters(
    dep_counters: Dict[str, Counter],
    sent_counters: Dict[str, Counter],
    normalize: bool
) -> Dict[str, pd.DataFrame]:
    out: Dict[str, pd.DataFrame] = {}

    def _counter_to_xy(counter: Counter, drop_non_positive: bool = False):
        items = sorted(counter.items(), key=lambda x: x[0])
        if drop_non_positive:
            items = [i for i in items if i[0] > 0]
        if normalize:
            total = sum(v for _, v in items)
            y = [v / total if total else 0 for _, v in items]
        else:
            y = [v for _, v in items]
        x = [k for k, _ in items]
        return x, y

    for m in ['dd', 'hd', 'v']:
        x, y = _counter_to_xy(dep_counters[m], drop_non_positive=True)
        out[m] = pd.DataFrame((x, y)).T

    for m in ['sl', 'tw', 'th', 'rd']:
        ctr = sent_counters[m]
        x, y = _counter_to_xy(ctr, drop_non_positive=False)
        out[m] = pd.DataFrame((x, y)).T

        for m, key in [('deprel', 'deprel'), ('pos', 'pos')]:
            items = sorted(dep_counters[key].items(), key=lambda x: x[1], reverse=True)
            if normalize:
                total = sum(v for _, v in items)
                y = [v / total if total else 0 for _, v in items]
            else:
                y = [v for _, v in items]
            x = [k for k, _ in items]
            out[m] = pd.DataFrame((x, y)).T
    return out


def _pvp_df_from_counters(dep_counter: Counter, gov_counter: Counter) -> pd.DataFrame:
    total_dep = sum(dep_counter.values())
    total_gov = sum(gov_counter.values())
    pvp = {
        'act as a dep': sorted(((k, v / total_dep if total_dep else 0) for k, v in dep_counter.items()), key=lambda x: x[1], reverse=True),
        'act as a gov': sorted(((k, v / total_gov if total_gov else 0) for k, v in gov_counter.items()), key=lambda x: x[1], reverse=True),
    }
    return _pvp2df(pvp)

def _build_sentence_ranges(sent_offsets: np.ndarray, workers: int, token_budget: int = 12000) -> List[Tuple[int, int]]:
    n_sent = len(sent_offsets) - 1
    if n_sent <= 0:
        return []
    ranges: List[Tuple[int, int]] = []
    cur_start = 0
    cur_tokens = 0
    for s in range(n_sent):
        sl = int(sent_offsets[s + 1] - sent_offsets[s])
        if cur_tokens >= token_budget and s > cur_start:
            ranges.append((cur_start, s))
            cur_start = s
            cur_tokens = 0
        cur_tokens += sl
    if cur_start < n_sent:
        ranges.append((cur_start, n_sent))
    # Keep enough tasks for load balancing.
    min_tasks = max(2 * workers, workers + 1)
    if len(ranges) < min_tasks and n_sent > min_tasks:
        step = max(1, n_sent // min_tasks)
        ranges = [(i, min(n_sent, i + step)) for i in range(0, n_sent, step)]
        if ranges and ranges[-1][1] < n_sent:
            ranges.append((ranges[-1][1], n_sent))
    return ranges

class DepValAnalyzer:

    def __init__(
        self,
        treebank,
        cleaned: bool = False,
        parallel: bool = False,
        n_jobs: Optional[int] = None,
        punct_deprels: Optional[List[str]] = None,
        root_deprels: Optional[List[str]] = None
    ):
        if type(treebank) is io.TextIOWrapper:
            self.raw = list(parse_incr(treebank))
        elif type(treebank) is list:
            self.raw = treebank
        elif type(treebank) is str:
            self.raw = parse(treebank)
        else:
            raise TypeError("treebank must be a text stream, list, or conllu string")
        self.parallel = parallel
        self.n_jobs = n_jobs
        self.punct_deprels = tuple(punct_deprels) if punct_deprels is not None else tuple(punctdeps)
        self.root_deprels = tuple(root_deprels) if root_deprels is not None else tuple(rootdeps)
        self.stop_deprels = tuple(set(self.punct_deprels).union({'_'}))
        if cleaned:
            self.treebank = treebank
        else:
            self.treebank = self.preprocessing()
        self.dep_metrics = ['dd','hd','ddir','v']
        self.sent_metrics = ['mdd','ndd','mhd','mhdd','tdl','sl','mv','vk','tw','th','hi','hf','rd']
        self.text_metrics = ['mdd','ndd','mhd','mhdd','mtdl','msl','mv','vk','mtw','mth','hi','hf','mrd']
        self.distribution_metrics = ['dd','hd','sl','v','tw','th','rd','deprel','pos']
        self.projection = {'mdd':'dd','mhd':'hd','mhdd':'hd','mtdl':'dd','msl':'id','mv':'v','vk':'v',
                           'mtw':'hd','mth':'hd','hi':'ddir','hf':'ddir','pos':'dpos','deprel':'deprel',
                           'rd':'dd','mrd':'dd','ndd':'dd','tw':'hd','th':'hd'}
        self._token_count = sum(len(sent) for sent in self.treebank)
        self._cache_dep: Dict[Tuple[Any, ...], Dict[str, List]] = {}
        self._cache_sent: Dict[Tuple[Any, ...], Dict[str, List]] = {}
        self._cache_text: Dict[Tuple[Any, ...], Dict[str, float]] = {}
        self._cache_dist: Dict[Tuple[Any, ...], Dict[str, Tuple[List, List]]] = {}
        self._cache_pvp: Dict[Tuple[Any, ...], Dict[str, List[Tuple[Any, float]]]] = {}

    def preprocessing(self):
        cleaned_treebank = []
        for sentence in self.raw:
            cleaned_sentence = []
            clean_id: Dict[int, int] = {}
            next_id = 1
            for word in sentence:
                wid = word.get('id')
                if isinstance(wid, int) and word.get('deprel') not in self.stop_deprels:
                    clean_id[wid] = next_id
                    next_id += 1
            for word in sentence:
                wid = word.get('id')
                if not isinstance(wid, int) or word.get('deprel') in self.stop_deprels:
                    continue
                new_word = dict(word)
                new_word['id'] = clean_id[wid]
                new_word['head'] = clean_id.get(word.get('head'), 0)
                cleaned_sentence.append(new_word)
            if len(cleaned_sentence) > 2:
                cleaned_treebank.append(cleaned_sentence)
        return cleaned_treebank  

    def _resolve_parallel(self, parallel: Optional[bool], n_jobs: Optional[int], n_items: int) -> Tuple[bool, int]:
        if parallel is None:
            parallel = self.parallel
        if n_jobs is None:
            n_jobs = self.n_jobs
        return _resolve_jobs(bool(parallel), n_jobs, n_items)

    def _resolve_parallel_runtime(self, parallel: Optional[bool], n_jobs: Optional[int], n_items: int) -> Tuple[bool, int]:
        use_parallel, workers = self._resolve_parallel(parallel, n_jobs, n_items)
        # Empirical gate: avoid multiprocessing when workload is too small for process overhead.
        min_tokens_per_worker = 120000
        if use_parallel and self._token_count < workers * min_tokens_per_worker:
            return False, 1
        return use_parallel, workers

    def _map_sentence_dep(
        self,
        metrics: List[str],
        parallel: Optional[bool],
        n_jobs: Optional[int],
        pool: Optional[Any] = None
    ) -> List[Dict[str, List[Any]]]:
        use_parallel, workers = self._resolve_parallel_runtime(parallel, n_jobs, len(self.treebank))
        payloads = [(sent, tuple(metrics), self.stop_deprels, self.root_deprels) for sent in self.treebank]
        if not use_parallel:
            return [_compute_sentence_dep(payload) for payload in payloads]
        sent_offsets = np.zeros(len(self.treebank) + 1, dtype=np.int64)
        if len(self.treebank) > 0:
            sent_offsets[1:] = np.cumsum([len(s) for s in self.treebank], dtype=np.int64)
        ranges = _build_sentence_ranges(sent_offsets, workers)
        temp_dir = tempfile.mkdtemp(prefix='depval_group_parts_')
        tasks = []
        for gid, (st, ed) in enumerate(ranges):
            tasks.append((gid, st, self.treebank[st:ed], tuple(metrics), self.stop_deprels, self.root_deprels, temp_dir))
        sentence_results: List[Optional[Dict[str, List[Any]]]] = [None] * len(self.treebank)
        try:
            if pool is not None:
                iterator = pool.imap_unordered(_sentence_group_worker_to_file, tasks, chunksize=1)
                for _, start_sent, part_path, _, _, _ in iterator:
                    with open(part_path, 'rb') as f:
                        part = pickle.load(f)
                    partial = _decode_group_results(part['results'], part['vocabs'])
                    for offset, sent_res in enumerate(partial):
                        sentence_results[start_sent + offset] = sent_res
            else:
                with _open_parallel_pool(processes=workers) as local_pool:
                    iterator = local_pool.imap_unordered(_sentence_group_worker_to_file, tasks, chunksize=1)
                    for _, start_sent, part_path, _, _, _ in iterator:
                        with open(part_path, 'rb') as f:
                            part = pickle.load(f)
                        partial = _decode_group_results(part['results'], part['vocabs'])
                        for offset, sent_res in enumerate(partial):
                            sentence_results[start_sent + offset] = sent_res
        finally:
            shutil.rmtree(temp_dir, ignore_errors=True)

        finalized: List[Dict[str, List[Any]]] = []
        for sent_res in sentence_results:
            if sent_res is None:
                raise RuntimeError("Sentence-level group parallelization produced incomplete results.")
            finalized.append(sent_res)
        return finalized

    def _calculate_sent_metrics_from_dep(self, dep_data: Dict[str, List], metrics: List[str]) -> Dict[str, List]:
        sent_data = {metric: [] for metric in metrics}
        for metric in metrics:
            if metric == 'mdd':
                sent_data[metric] = [sum(i for i in j if i > 0) / max(1, len(j) - 1) for j in dep_data['dd']]
            elif metric == 'mhd':
                sent_data[metric] = [sum(i for i in j if i > 0) / max(1, len(j) - 1) for j in dep_data['hd']]
            elif metric == 'mhdd':
                sent_data[metric] = [(len(j) - 1) / (max(j) + 1) if j else 0 for j in dep_data['hd']]
            elif metric == 'tdl':
                sent_data[metric] = [sum(j) for j in dep_data['dd']]
            elif metric == 'sl':
                sent_data[metric] = [len(j) for j in dep_data['id']]
            elif metric == 'mv':
                sent_data[metric] = [sum(j) / max(1, len(j)) for j in dep_data['v']]
            elif metric == 'vk':
                sent_data[metric] = [float(np.var(j)) for j in dep_data['v']]
            elif metric == 'tw':
                sent_data[metric] = [_safe_twig(j) for j in dep_data['hd']]
            elif metric == 'th':
                sent_data[metric] = [max(j) + 1 if j else 0 for j in dep_data['hd']]
            elif metric == 'hi':
                sent_data[metric] = [j.count(1) for j in dep_data['ddir']]
            elif metric == 'hf':
                sent_data[metric] = [j.count(-1) for j in dep_data['ddir']]
            elif metric == 'rd':
                sent_data[metric] = [_safe_root_distance(j) for j in dep_data['dd']]
            elif metric == 'ndd':
                values = []
                for j in dep_data['dd']:
                    mdd = sum(i for i in j if i > 0) / max(1, len(j) - 1)
                    rd = _safe_root_distance(j)
                    denom = np.sqrt(max(1, rd) * max(1, len(j)))
                    values.append(abs(np.log(max(mdd, 1e-12) / max(denom, 1e-12))))
                sent_data[metric] = values
        return sent_data

    def _calculate_text_metrics_from_dep(self, dep_data: Dict[str, List], metrics: List[str]) -> Dict[str, float]:
        text_data: Dict[str, float] = {}
        flat_dd = _flatten(dep_data.get('dd', []))
        flat_hd = _flatten(dep_data.get('hd', []))
        flat_v = _flatten(dep_data.get('v', []))
        flat_ddir = _flatten(dep_data.get('ddir', []))
        n_sent = max(1, len(dep_data.get('dd', [])))

        for metric in metrics:
            if metric == 'mdd':
                text_data[metric] = sum(flat_dd) / max(1, len(flat_dd) - len(dep_data['dd']))
            elif metric == 'mhd':
                text_data[metric] = sum(flat_hd) / max(1, len(flat_hd) - len(dep_data['hd']))
            elif metric == 'mhdd':
                text_data[metric] = sum((len(j) - 1) / (max(j) + 1) if j else 0 for j in dep_data['hd']) / n_sent
            elif metric == 'mtdl':
                text_data[metric] = sum(sum(j) for j in dep_data['dd']) / n_sent
            elif metric == 'msl':
                text_data[metric] = sum(len(j) for j in dep_data['dd']) / n_sent
            elif metric == 'mv':
                text_data[metric] = sum(flat_v) / max(1, len(flat_v))
            elif metric == 'vk':
                text_data[metric] = sum(float(np.var(j)) for j in dep_data['v']) / max(1, len(dep_data['v']))
            elif metric == 'mtw':
                text_data[metric] = sum(_safe_twig(j) for j in dep_data['hd']) / max(1, len(dep_data['hd']))
            elif metric == 'mth':
                text_data[metric] = sum((max(j) + 1) if j else 0 for j in dep_data['hd']) / max(1, len(dep_data['hd']))
            elif metric == 'hi':
                text_data[metric] = sum(j.count(1) for j in dep_data['ddir']) / max(1, len(flat_ddir) - len(dep_data['ddir']))
            elif metric == 'hf':
                text_data[metric] = sum(j.count(-1) for j in dep_data['ddir']) / max(1, len(flat_ddir) - len(dep_data['ddir']))
            elif metric == 'mrd':
                text_data[metric] = sum(_safe_root_distance(j) for j in dep_data['dd']) / n_sent
            elif metric == 'ndd':
                sent_ndd = self._calculate_sent_metrics_from_dep(dep_data, ['ndd'])['ndd']
                text_data[metric] = sum(sent_ndd) / max(1, len(sent_ndd))
        return text_data

    def _calculate_distributions_from_dep(
        self,
        dep_data: Dict[str, List],
        metrics: List[str],
        normalize: bool
    ) -> Dict[str, Tuple[List, List]]:
        distributions: Dict[str, Tuple[List, List]] = {}

        def normalize_data(data_items):
            if not normalize:
                return [i[1] for i in data_items]
            total = sum(i[1] for i in data_items)
            if total == 0:
                return [0 for _ in data_items]
            return [i[1] / total for i in data_items]

        metric_functions = {
            'sl': lambda: [len(j) for j in dep_data['id']],
            'tw': lambda: [_safe_twig(j) for j in dep_data['hd']],
            'th': lambda: [(max(j) + 1) if j else 0 for j in dep_data['hd']],
            'rd': lambda: [_safe_root_distance(j) for j in dep_data['dd']],
        }

        for metric in metrics:
            if metric in metric_functions:
                data = sorted(Counter(metric_functions[metric]()).items())
                x = [i[0] for i in data]
                y = normalize_data(data)
                distributions[metric] = (x, y)
                continue

            if metric in ['deprel', 'pos']:
                source_metric = self.projection.get(metric, metric)
                data = sorted(
                    Counter(_flatten(dep_data.get(source_metric, []))).items(),
                    key=lambda x: x[1],
                    reverse=True
                )
                x = [i[0] for i in data]
                y = normalize_data(data)
                distributions[metric] = (x, y)
                continue

            flattened = _flatten(dep_data.get(metric, []))
            data = sorted(Counter(flattened).items())
            data = [i for i in data if i[0] > 0]
            x = [i[0] for i in data]
            y = normalize_data(data)
            distributions[metric] = (x, y)

        return distributions

    def calculate_dep_metrics(
        self,
        metrics: List[str] = None,
        parallel: Optional[bool] = None,
        n_jobs: Optional[int] = None
    ) -> Dict[str, List]:
        if metrics is None:
            metrics = self.dep_metrics
        base_metrics = ['id', 'form', 'dpos', 'head', 'gform', 'gpos', 'deprel']
        metrics = base_metrics + [m for m in metrics if m not in base_metrics]
        if parallel is None:
            parallel = self.parallel
        if n_jobs is None:
            n_jobs = self.n_jobs
        pkey = self._resolve_parallel_runtime(parallel, n_jobs, len(self.treebank))
        cache_key = (tuple(metrics), pkey)
        if cache_key in self._cache_dep:
            return self._cache_dep[cache_key]
        sentence_results = self._map_sentence_dep(metrics, parallel=parallel, n_jobs=n_jobs)
        dep_data = {metric: [sent[metric] for sent in sentence_results] for metric in metrics}
        self._cache_dep[cache_key] = dep_data
        return dep_data
    
    def calculate_sent_metrics(
        self,
        metrics: List[str] = None,
        parallel: Optional[bool] = None,
        n_jobs: Optional[int] = None
    ) -> Dict[str, List]:
        if metrics is None:
            metrics = self.sent_metrics
        if parallel is None:
            parallel = self.parallel
        if n_jobs is None:
            n_jobs = self.n_jobs
        pkey = self._resolve_parallel_runtime(parallel, n_jobs, len(self.treebank))
        cache_key = (tuple(metrics), pkey)
        if cache_key in self._cache_sent:
            return self._cache_sent[cache_key]
        dep_metrics = list(set([self.projection[metric] for metric in metrics if metric in self.projection]))
        dep_data = self.calculate_dep_metrics(dep_metrics, parallel=parallel, n_jobs=n_jobs)
        sent_data = self._calculate_sent_metrics_from_dep(dep_data, metrics)
        self._cache_sent[cache_key] = sent_data
        return sent_data

    
    def calculate_text_metrics(
        self,
        metrics: List[str] = None,
        parallel: Optional[bool] = None,
        n_jobs: Optional[int] = None
    ) -> Dict[str, float]:
        if metrics is None:
            metrics = self.text_metrics
        if parallel is None:
            parallel = self.parallel
        if n_jobs is None:
            n_jobs = self.n_jobs
        pkey = self._resolve_parallel_runtime(parallel, n_jobs, len(self.treebank))
        cache_key = (tuple(metrics), pkey)
        if cache_key in self._cache_text:
            return self._cache_text[cache_key]
        dep_metrics = list(set([self.projection[metric] for metric in metrics if metric in self.projection]))
        dep_data = self.calculate_dep_metrics(dep_metrics, parallel=parallel, n_jobs=n_jobs)
        text_data = self._calculate_text_metrics_from_dep(dep_data, metrics)
        self._cache_text[cache_key] = text_data
        return text_data

    def calculate_distributions(
        self,
        metrics: List[str] = None,
        normalize: bool = False,
        parallel: Optional[bool] = None,
        n_jobs: Optional[int] = None
    ) -> Dict[str, Tuple[List, List]]:
        if metrics is None:
            metrics = self.distribution_metrics
        if parallel is None:
            parallel = self.parallel
        if n_jobs is None:
            n_jobs = self.n_jobs
        pkey = self._resolve_parallel_runtime(parallel, n_jobs, len(self.treebank))
        cache_key = (tuple(metrics), normalize, pkey)
        if cache_key in self._cache_dist:
            return self._cache_dist[cache_key]
        dep_metrics = list(set(self.projection.get(metric, metric) for metric in metrics))
        dep_data = self.calculate_dep_metrics(dep_metrics, parallel=parallel, n_jobs=n_jobs)
        dist_data = self._calculate_distributions_from_dep(dep_data, metrics, normalize)
        self._cache_dist[cache_key] = dist_data
        return dist_data

    def calculate_pvp(
        self,
        input: Optional[str] = None,
        target: str = 'deprel',
        normalize: bool = True,
        parallel: Optional[bool] = None,
        n_jobs: Optional[int] = None
    ):
        payloads = [(sentence, input, target, self.stop_deprels) for sentence in self.treebank]
        if parallel is None:
            parallel = self.parallel
        if n_jobs is None:
            n_jobs = self.n_jobs
        pkey = self._resolve_parallel_runtime(parallel, n_jobs, len(payloads))
        cache_key = (input, target, normalize, pkey)
        if cache_key in self._cache_pvp:
            return self._cache_pvp[cache_key]

        # Fast path: if dep-level base vectors are already cached, derive pvp directly.
        dep_metrics = ['id', 'head', 'dpos', 'deprel']
        dep_cache_key = (tuple(dep_metrics), pkey)
        dep_data_cached = self._cache_dep.get(dep_cache_key)
        if dep_data_cached is not None:
            pvp_data = _calculate_pvp_from_dep_data(
                dep_data_cached,
                input_upos=input,
                target=target,
                normalize=normalize
            )
            self._cache_pvp[cache_key] = pvp_data
            return pvp_data

        use_parallel, workers = self._resolve_parallel_runtime(parallel, n_jobs, len(payloads))
        if use_parallel:
            chunksize = _recommended_chunksize(len(payloads), workers)
            with _open_parallel_pool(processes=workers) as pool:
                sentence_counts = pool.map(_compute_sentence_pvp, payloads, chunksize=chunksize)
        else:
            sentence_counts = [_compute_sentence_pvp(payload) for payload in payloads]

        dependents = Counter()
        governors = Counter()
        for dep_part, gov_part in sentence_counts:
            dependents.update(dep_part)
            governors.update(gov_part)

        total_deps = sum(dependents.values())
        total_govs = sum(governors.values())
        if normalize:
            govs = sorted(
                ((k, (v / total_govs) if total_govs else 0) for k, v in governors.items()),
                key=lambda x: x[1],
                reverse=True
            )
            deps = sorted(
                ((k, (v / total_deps) if total_deps else 0) for k, v in dependents.items()),
                key=lambda x: x[1],
                reverse=True
            )
        else:
            govs = sorted(governors.items(), key=lambda x: x[1], reverse=True)
            deps = sorted(dependents.items(), key=lambda x: x[1], reverse=True)
        pvp_data = {'act as a gov': govs, 'act as a dep': deps}
        self._cache_pvp[cache_key] = pvp_data
        return pvp_data

    def _compute_feature_bundle(
        self,
        normalize: bool = True,
        parallel: Optional[bool] = None,
        n_jobs: Optional[int] = None,
        compute_mode: str = 'full-output'
    ) -> Tuple[Dict[str, float], pd.DataFrame, pd.DataFrame, pd.DataFrame, Dict[str, pd.DataFrame]]:
        if parallel is None:
            parallel = self.parallel
        if n_jobs is None:
            n_jobs = self.n_jobs
        use_parallel, workers = self._resolve_parallel_runtime(parallel, n_jobs, len(self.treebank))
        dep_metrics = ['id', 'dpos', 'head', 'deprel'] + list(self.dep_metrics)
        if compute_mode == 'full-output':
            dep_metrics = ['id', 'form', 'dpos', 'head', 'gform', 'gpos', 'deprel'] + list(self.dep_metrics)
        dep_metrics = list(dict.fromkeys(dep_metrics))
        dep_data = self.calculate_dep_metrics(
            metrics=dep_metrics,
            parallel=use_parallel,
            n_jobs=workers if use_parallel else 1
        )

        # Derive default pvp directly from dep vectors to avoid an extra full-treebank pass.
        pvp_data = _calculate_pvp_from_dep_data(dep_data, input_upos=None, target='deprel', normalize=True)

        sent_data = self._calculate_sent_metrics_from_dep(dep_data, self.sent_metrics)
        text_data = self._calculate_text_metrics_from_dep(dep_data, self.text_metrics)
        distributions = self._calculate_distributions_from_dep(dep_data, self.distribution_metrics, normalize=normalize)

        dep_frame = pd.DataFrame({x: _flatten(y) for x, y in dep_data.items()}) if compute_mode == 'full-output' else pd.DataFrame()
        sent_frame = pd.DataFrame(sent_data)
        pvp_frame = _pvp2df(pvp_data)
        distribution_dict = {m: pd.DataFrame(distributions[m]).T for m in self.distribution_metrics}
        return text_data, dep_frame, sent_frame, pvp_frame, distribution_dict
    
def _pvp2df(pvp):
    df = pd.DataFrame()
    gov = dict(pvp['act as a gov'])
    dep = dict(pvp['act as a dep'])
    df['Items'] = sorted(list(set(gov.keys()).union(set(dep.keys()))))  
    df['act as a gov'] = [gov.get(d,0) for d in df['Items']]
    df['act as a dep'] = [dep.get(d,0) for d in df['Items']]
    return df

def getDepValFeatures(
    treebank: list,
    normalize: bool = True,
    parallel: bool = False,
    n_jobs: Optional[int] = None,
    punct_deprels: Optional[List[str]] = None,
    root_deprels: Optional[List[str]] = None
):
    analyzer = DepValAnalyzer(
        treebank,
        parallel=parallel,
        n_jobs=n_jobs,
        punct_deprels=punct_deprels,
        root_deprels=root_deprels
    )
    return analyzer._compute_feature_bundle(normalize=normalize, parallel=parallel, n_jobs=n_jobs)


def analyze(
    treebank_path: str,
    out_path: str,
    normalize: bool = True,
    parallel: bool = False,
    n_jobs: Optional[int] = None,
    punct_deprels: Optional[List[str]] = None,
    root_deprels: Optional[List[str]] = None
):

    if not os.path.splitext(out_path)[1]:
        os.makedirs(out_path,exist_ok=True)
    else:
        raise ValueError(f" The output path {out_path} is not a valid directory. It should be a directory. ")

    if os.path.isdir(treebank_path):
        treebanks = [os.path.join(treebank_path, i) for i in os.listdir(treebank_path) if i.endswith('.conllu')]
        treebanks.sort()
        text_csv = pd.DataFrame(columns=['mdd','ndd','mhd','mhdd','mtdl','msl','mv','vk','mtw','mth','hi','hf','mrd'])
        file_names = [os.path.basename(t).split('.')[0] for t in treebanks]
        files = [os.path.join(out_path, f) for f in file_names]
        file_stats = []
        for i, t in enumerate(treebanks):
            file_stats.append({
                'index': i,
                'treebank_path': t,
                'file_name': file_names[i],
                'out_dir': files[i],
                'file_size': os.path.getsize(t),
                'sentence_count': _estimate_sentence_count(t),
            })

        mode, workers = _select_parallel_mode(parallel=parallel, n_jobs=n_jobs, file_stats=file_stats)
        if mode == 'file-level':
            tasks = []
            for stat in sorted(file_stats, key=lambda x: x['sentence_count'], reverse=True):
                os.makedirs(stat['out_dir'], exist_ok=True)
                task = {
                    **stat,
                    'normalize': normalize,
                    'punct_deprels': punct_deprels,
                    'root_deprels': root_deprels,
                }
                tasks.append(task)
            batch_size = max(1, len(tasks) // max(1, workers * 2))
            task_batches = _split_batches(tasks, batch_size)
            with _open_parallel_pool(processes=workers) as pool:
                for batch_results in pool.imap_unordered(_analyze_file_batch_worker, task_batches, chunksize=1):
                    for idx, file_name, text_metrics in batch_results:
                        text_csv.loc[idx, :] = text_metrics
        else:
            sentence_parallel = (mode == 'sentence-level')
            sentence_jobs = workers if sentence_parallel else 1
            for i,t in enumerate(treebanks):      
                os.makedirs(files[i],exist_ok=True)     
                if not sentence_parallel:
                    text_metrics = _analyze_single_file_serial(
                        treebank_path=t,
                        out_dir=files[i],
                        normalize=normalize,
                        punct_deprels=punct_deprels,
                        root_deprels=root_deprels
                    )
                else:
                    text_metrics = _analyze_single_file_sentence_level(
                        treebank_path=t,
                        out_dir=files[i],
                        normalize=normalize,
                        workers=sentence_jobs,
                        punct_deprels=punct_deprels,
                        root_deprels=root_deprels
                    )
                text_csv.loc[i,:] = text_metrics
        
        text_csv = text_csv.reindex(range(len(file_names)))
        text_csv = text_csv.astype(float).round(2)
        text_csv['treebank'] = [file_names[i] for i in range(len(file_names))]
        text_csv.to_csv(os.path.join(out_path, f'text_metrics.csv'), index=False)
    elif os.path.isfile(treebank_path):
        file_name = os.path.basename(treebank_path).split('.')[0]
        file_stats = [{
            'index': 0,
            'treebank_path': treebank_path,
            'file_name': file_name,
            'out_dir': out_path,
            'file_size': os.path.getsize(treebank_path),
            'sentence_count': _estimate_sentence_count(treebank_path),
        }]
        mode, workers = _select_parallel_mode(parallel=parallel, n_jobs=n_jobs, file_stats=file_stats)
        if mode != 'sentence-level':
            text_metrics = _analyze_single_file_serial(
                treebank_path=treebank_path,
                out_dir=out_path,
                normalize=normalize,
                punct_deprels=punct_deprels,
                root_deprels=root_deprels
            )
        else:
            text_metrics = _analyze_single_file_sentence_level(
                treebank_path=treebank_path,
                out_dir=out_path,
                normalize=normalize,
                workers=workers,
                punct_deprels=punct_deprels,
                root_deprels=root_deprels
            )
        text_csv = pd.DataFrame([text_metrics])
        text_csv = text_csv.astype(float).round(2)
        text_csv['treebank'] = [file_name]
        text_csv.to_csv(os.path.join(out_path, 'text_metrics.csv'), index=False)
    else:
        raise ValueError(f"treebank_path {treebank_path} is neither a directory nor a file.")

class Converter:
     
    def __init__(self, treebank):
        self.treebank = treebank

    def to_conllu(self,style:str):

        if style == 'pmt':
            treebank = self.treebank.read()
            sents = treebank.strip().split('\n\n')
            conllu_sents = []
            for sent in sents:
                columns = [[i for i in line.split('\t') if i!=''] for line in sent.split('\n') if line]

                structured_sent = [
                    {
                        'id': i + 1,
                        'form': columns[0][i],
                        'lemma': '_',
                        'upos': columns[1][i],
                        'xpos': '_',
                        'feats': '_',
                        'head': columns[3][i],
                        'deprel': columns[2][i],
                        'deps': '_',
                        'misc': '_'
                    }
                    for i in range(len(columns[0])) if columns[3][i] != '_'
                ]
                
                conllu_sents.append(structured_sent)
            return conllu_sents
        
        elif style == 'conll':
            sents = parse_incr(self.treebank)
            conllu_sents = [[{**w, 'deps': '_','misc': '_'} for w in s] for s in sents]
            return conllu_sents
        
        elif style == 'mcdt':
            treebank = pd.read_csv(self.treebank, delimiter='\t')
            treebank = treebank.dropna(subset=[treebank.columns[4]])

            def senter(df: pd.DataFrame):
                index_slices = []
                current_indices = []

                for i in range(len(df)):
                    
                    if not current_indices or df.iloc[i, 1] > df.iloc[current_indices[-1], 1]:
                        current_indices.append(i)
                    else:
                        index_slices.append(current_indices.copy())
                        current_indices = [i]

                if current_indices:
                    index_slices.append(current_indices)
                
                return index_slices
            
            sent_ids = senter(treebank)
            conllu_sents = [
                [
                    {
                        'id': int(treebank.iloc[i, 1]),
                        'form': treebank.iloc[i, 2],
                        'lemma': '_',
                        'upos': treebank.iloc[i, 3],
                        'xpos': '_',
                        'feats': '_',
                        'head':  0 if treebank.iloc[i, 7] == 's' else int(treebank.iloc[i, 4]),
                        'deprel': treebank.iloc[i, 7],
                        'deps': '_',
                        'misc': '_'
                    }
                    for i in s
                ]
                for s in sent_ids]

            return conllu_sents
    
    def to_others(self,style:str,cache=None):
        if cache is None:
            sents = parse_incr(self.treebank)
        else:
            sents = cache
        if style == 'conll':
            conll_sents = [[{k: v for k, v in w.items() if k not in ['deps','misc']} for w in s] for s in sents]
            return conll_sents
        
        elif style == 'pmt':
            pmt_sents = [[
                [w['form'] for w in s],
                [w['upos'] for w in s],
                [w['deprel'] for w in s],
                [w['head'] for w in s]
            ] for s in sents]
            return pmt_sents
        
        elif style == 'mcdt':
            mcdt_sents = pd.DataFrame(columns=['sent','id','form', 'pos', 'head', 'gform','gpos','deprel'])
            sent_id = 0
            for s in sents:
                sent_id += 1
                for w in s:
                    if w['head'] != '_' and w['head'] is not None:
                        mcdt_sents.loc[len(mcdt_sents)] = [sent_id, w['id'], w['form'], w['upos'], w['head'], s[int(w['head'])-1]['form'], s[int(w['head'])-1]['upos'], w['deprel']]
            return mcdt_sents.to_csv(sep='\t', index=False).replace('\r','')
        
    def style2style(self,style_from:str,style_to:str):
        if style_from == style_to:
            raise ValueError(f"The input and output styles should be different. Got {style_from} and {style_to}.")
        elif style_from == 'conllu' and style_to != 'conllu':
            treebank = self.to_others(style_to)
        elif style_from != 'conllu' and style_to == 'conllu':
            treebank = self.to_conllu(style_from)
        elif style_from != 'conllu' and style_to != 'conllu':
            cache = self.to_conllu(style_from)
            treebank = self.to_others(style_to,cache=cache)
        return treebank
    
    def save(self,treebank,style:str,file_path:str):
        def format_feats(feats):
            if feats and feats != '_':
                return '|'.join(f"{k}={v}" for k, v in feats.items())
            return '_'
        
        if style == 'conllu':
            with open(file_path, 'w', encoding='utf-8') as f:
                for s in treebank:
                    for w in s:
                        if isinstance(w['id'], int) or isinstance(w['id'], str): 
                            f.write(f"{w['id']}\t{w['form']}\t{w['lemma']}\t{w['upos']}\t{w['xpos']}\t{format_feats(w['feats'])}\t{w['head']}\t{w['deprel']}\t{w['deps']}\t{w['misc']}\n")
                        elif isinstance(w['id'], tuple) :
                            f.write(f"{w['id'][0]}{w['id'][1]}{w['id'][2]}\t{w['form']}\t_\t_\t_\t_\t_\t_\t_\t_\n")
                    f.write('\n')
        elif style == 'pmt':
            with open(file_path, 'w', encoding='utf-8') as f:
                for s in treebank:
                    f.write('\t'.join(s[0]) + '\n')
                    f.write('\t'.join(s[1]) + '\n')
                    f.write('\t'.join(s[2]) + '\n')
                    f.write('\t'.join(map(lambda x: str(x) if x is not None else '_', s[3])))
                    f.write('\n\n')
        elif style == 'conll':
            with open(file_path, 'w', encoding='utf-8') as f:
                for s in treebank:
                    for w in s:
                        if isinstance(w['id'], int) or isinstance(w['id'], str): 
                            f.write(f"{w['id']}\t{w['form']}\t{w['lemma']}\t{w['upos']}\t{w['xpos']}\t{format_feats(w['feats'])}\t{w['head']}\t{w['deprel']}\n")
                        elif isinstance(w['id'], tuple):
                            f.write(f"{w['id'][0]}{w['id'][1]}{w['id'][2]}\t{w['form']}\t_\t_\t_\t_\t_\t_\n")
                    f.write('\n')
        elif style == 'mcdt':
            with open(file_path, 'w', encoding='utf-8') as f:
                f.write(treebank)
        else:
            raise ValueError(f"Invalid style {style}.")

def convert(treebank_path:str,out_path:str,style_from:str,style_to:str):

    if not os.path.splitext(out_path)[1]:
        os.makedirs(out_path,exist_ok=True)
    else:    
        raise ValueError(f" The output path {out_path} is not a valid directory. It should be a directory. ")
    
    if os.path.isdir(treebank_path):
            treebanks = [os.path.join(treebank_path, i) for i in os.listdir(treebank_path)]
            file_names = [os.path.basename(t).split('.')[0] for t in treebanks]
            for i,t in enumerate(treebanks):
                with open(t, encoding='utf-8') as treebank:
                    converter = Converter(treebank)
                    converted_treebank = converter.style2style(style_from,style_to)
                    converter.save(converted_treebank,style_to,os.path.join(out_path,f'{file_names[i]}.conllu'))
    else:
        raise ValueError(f"The input path {treebank_path} is not a valid directory. It should be a directory containing treebank files.")




