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

import pandas as pd
from conllu import parse_incr,parse
from collections import Counter
import os
import io
import random
import numpy as np
from collections import Counter,defaultdict,deque
from typing import List, Dict, Optional, Tuple, Any
import multiprocessing as mp
from tqdm import tqdm
import time
import networkx as nx

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

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

class DepValAnalyzer:

    def __init__(self, treebank,cleaned=False):
        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)
        if cleaned:
            self.treebank = treebank
        elif not cleaned:
            self.treebank = self.preprocessing()
        self.dep_metrics = ['dd','hd','ddir','v','l']
        self.sent_metrics = ['mdd','ndd','mhd','mhdd','tdl','sl','mv','vk','vl','tw','th','hi','hf','rd','mndl']
        self.text_metrics = ['mdd','ndd','mhd','mhdd','mtdl','msl','mv','vk','vl','mtw','mth','hi','hf','mrd','mndl']
        self.distribution_metrics = ['dd','hd','sl','v','tw','th','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',
                           'vl':'l','rd':'dd','mrd':'dd','ndd':'dd','tdl':'dd','mndl':'hd','phf':'ddir','rmdd':'ddir','lmdd':'ddir',
                           'rmhd':'ddir','lmhd':'ddir','sdd':'dd','shd':'dd','lsdd':'dd','rsdd':'dd',
                           'lshd':'hd','rshd':'hd','sldc':'ddir','srdc':'ddir'}

    def preprocessing(self):
        cleaned_treebank = []
        for sentence in self.raw:
            cleaned_sentence = []
            punkt_id = [word['id'] for word in sentence if word['deprel'] in punctdeps]
            clean_id = {word['id']:word['id']-len([i for i in punkt_id if i < word['id']]) for word in sentence if word['deprel'] not in stopdeps}
            for word in sentence:
                if word['deprel'] not in stopdeps:
                    word['id'] = clean_id[word['id']]
                    word['head'] = clean_id.get(word['head'],0)
                    cleaned_sentence.append(word)
            if len(cleaned_sentence) > 2:
                cleaned_treebank.append(cleaned_sentence)
        return cleaned_treebank  

    def _dd(self,word: Dict[str, Any])->int:
        dd = abs(word['head'] - word['id'])
        return dd
     
    def _hd(self,word: Dict[str, Any],word_index_by_id: Dict[int, Dict[str, Any]])->int:
        head_id = word['head']
        hd = 0
        while head_id != 0:
            hd += 1
            if head_id not in word_index_by_id.keys():
                hd = 0
                break
            elif head_id in word_index_by_id.keys():
                head_id = word_index_by_id[head_id]['head']
            if hd > len(word_index_by_id):
                hd = 0
                break
        return hd
    
    def _ddir(self,word: Dict[str, Any])->int:
        if word['head'] < word['id']:
            return 1
        elif word['head'] > word['id']:
            return -1

    def _v(self, word: Dict[str, Any],sent: List[Dict[str, Any]]) -> int:
        v = sum(1 for w in sent if w['head'] == word['id'] and w['deprel'] not in stopdeps) + 1
        if word['head'] == 0:
            v -= 1
        return v
    
    # def _l(self, word: Dict[str, Any],sent: List[Dict[str, Any]]) -> int:
    #     edges = [(w['id']-1,w['head']-1) for w in sent if w['head'] != 0]
    #     src = word['id']-1
    #     V = len(edges) + 1
    #     adj = [[] for _ in range(V)]
    #     for u, v in edges:
    #         adj[u].append(v)
    #         adj[v].append(u)
    #     dist = [-1] * V
    #     dist[src] = 0
    #     q = deque([src])
    #     while q:
    #         u = q.popleft()
    #         for v in adj[u]:
    #             if dist[v] == -1:
    #                 dist[v] = dist[u] + 1
    #                 q.append(v)
    #     total, count = sum(d for d in dist if d > 0), sum(1 for d in dist if d > 0)
    #     return total / max(1, count)
    def _l(self, word: Dict[str, Any], sent: List[Dict[str, Any]]) -> int:
        try:
            edges = [(w['id'] - 1, w['head'] - 1) for w in sent if w['head'] != 0]
            src = word['id'] - 1
            V = len(edges) + 1
            adj = [[] for _ in range(V)]

            for u, v in edges:
                if u >= V or v >= V or u < 0 or v < 0:
                    print("❌ IndexError: 节点超出范围")
                    print(f"出错句子: {sent}")
                    import sys
                    sys.exit(1)
                adj[u].append(v)
                adj[v].append(u)

            dist = [-1] * V
            dist[src] = 0
            q = deque([src])
            while q:
                u = q.popleft()
                for v in adj[u]:
                    if dist[v] == -1:
                        dist[v] = dist[u] + 1
                        q.append(v)

            total = sum(d for d in dist if d > 0)
            count = sum(1 for d in dist if d > 0)
            return total / max(1, count)

        except Exception as e:
            print(f"❌ 计算 _l() 时出错: {type(e).__name__} — {e}")
            print(f"出错句子: {sent}")
            import sys
            sys.exit(1)

    def _vl(self,heads, directed=False):
        edges = [(i,heads[i]-1) for i in range(len(heads))]
        edges = [e for e in edges if e[1] >= 0]
        g = defaultdict(list)
        for u, v in edges:
            g[u].append(v)
            if not directed:
                g[v].append(u)
        nodes = sorted(g)
        lengths = []
        for i, s in enumerate(nodes):
            dist = {s: 0}
            q = deque([s])
            while q:
                u = q.popleft()
                for v in g[u]:
                    if v not in dist:
                        dist[v] = dist[u] + 1
                        q.append(v)
            for t in nodes[i+1:]:
                if t in dist:
                    lengths.append(dist[t])
        return np.var(lengths)
    
    def edit_distance_from_baseline(self,heads,baseline='chain'):
        def generate_chains(sl=10,root=None):
            if root is None:
                root = random.choice([1,sl])
            left_links = [i+1 for i in range(1,root)]
            right_links = [i-1 for i in range(root+1,sl+1)]
            links = left_links + [0] + right_links
            return links    
        def generate_stars(sl=10,root=None):
            if root is None:
                root = random.choice([1,sl])
            links = [0 if i == root else root for i in range(1, sl + 1)]
            return links
        def generate_nests(sl=10,root=None):
            if root is None:
                root = random.choice([1, sl])
            links = []
            l, r = (1, sl) if root == 1 else (sl, 1)
            step = 1 if root == 1 else -1

            for i in range(sl - 1):
                links.append((l, r))
                l, r = r, l + step if i % 2 == 0 else l - step
            links = links if root == 1 else [(i[1],i[0]) for i in links]
            links.append((root, 0))
            links = sorted(links)
            sent = [i[1] for i in links]
            return sent
        root = heads.index(0)+1
        if baseline=='chain':
            # baseline1 = generate_chains(len(heads),1)
            # baseline2 = generate_chains(len(heads),len(heads))
            baseline = generate_chains(len(heads),root)
            ed = [heads[i]!= baseline[i] for i in range(len(heads))].count(True)
        elif baseline=='star':
            # baseline1 = generate_stars(len(heads),1)
            # baseline2 = generate_stars(len(heads),len(heads))
            baseline = generate_stars(len(heads),root)
            ed = [heads[i]!= baseline[i] for i in range(len(heads))].count(True)
        elif baseline=='nest':
            baseline1 = generate_nests(len(heads),1)
            baseline2 = generate_nests(len(heads),len(heads))
            ed1 = [heads[i]!= baseline1[i] for i in range(len(heads))].count(True)
            ed2 = [heads[i]!= baseline2[i] for i in range(len(heads))].count(True)
            ed =  min(ed1,ed2)

        # ed = [heads[i]!= baseline[i] for i in range(len(heads))].count(True)
        return ed
    def dc(self,ddirs):
        root = ddirs.index(0)
        left = ddirs[:root]
        ldc = left.count(-1)/len(left)
        right = ddirs[root+1:]
        rdc = right.count(1)/len(ddirs[root+1:])
        return ldc,rdc


    def calculate_dep_metrics(self, metrics: List[str] = None) -> Dict[str, List]:
        if metrics is None:
            metrics = self.dep_metrics
        metrics = ['sent_id','id','form','dpos','head','gform','gpos','deprel']+[m for m in metrics if m not in ['id','dpos','head','gpos','deprel']]
        results = {metric: [] for metric in metrics}
        
        sent_id = 0
        for sent in self.treebank:
            word_index_by_id = {word['id']: word for word in sent if type(word['id']) != tuple } if 'hd' in metrics or 'dd' in metrics else None
            temp_results = {metric: [] for metric in metrics}
            
            for word in sent:
                if word['head'] != 0:
                    if 'dd' in metrics:
                        temp_results['dd'].append(self._dd(word))  
                    if 'hd' in metrics:
                        temp_results['hd'].append(self._hd(word, word_index_by_id))
                    if 'ddir' in metrics:
                        temp_results['ddir'].append(self._ddir(word)) 
                    if 'v' in metrics:
                        temp_results['v'].append(self._v(word, sent)) 
                    if 'l' in metrics:
                        temp_results['l'].append(self._l(word, sent))
                    temp_results['id'].append(word['id'])
                    temp_results['form'].append(word['form'])
                    temp_results['dpos'].append(word['upos'])
                    temp_results['head'].append(word['head'])
                    try:
                        temp_results['gform'].append(sent[word['head']-1]['form'])
                    except:
                        temp_results['gform'].append('ERROR_ANNOTATION')
                    try:
                        temp_results['gpos'].append(sent[word['head']-1]['upos'])
                    except:
                        temp_results['gpos'].append('ERROR_ANNOTATION')
                    temp_results['deprel'].append(word['deprel'])
                elif word['head'] == 0:
                    if 'dd' in metrics:
                        temp_results['dd'].append(0)  
                    if 'hd' in metrics:
                        temp_results['hd'].append(0)
                    if 'ddir' in metrics:
                        temp_results['ddir'].append(0) 
                    if 'v' in metrics:
                        temp_results['v'].append(self._v(word, sent)) 
                    if 'l' in metrics:
                        temp_results['l'].append(self._l(word, sent))
                    temp_results['id'].append(word['id'])
                    temp_results['form'].append(word['form'])
                    temp_results['dpos'].append(word['upos'])
                    temp_results['head'].append(word['head'])
                    temp_results['gform'].append('ROOT')
                    temp_results['gpos'].append('ROOT')
                    temp_results['deprel'].append(word['deprel'])
                temp_results['sent_id'].append(sent_id)
            sent_id += 1

            for metric in results:
                results[metric].append(temp_results[metric])

        return results
    
    def dep_info(self, sent: List[List[Dict[str, Any]]], metrics: List[str] = None,sent_id=0) -> Dict[str, List]:
        if metrics is None:
            metrics = self.dep_metrics

        dep_metrics = metrics + ['id','sent_id']
        word_index_by_id = {word['id']: word for word in sent if type(word['id']) != tuple } if 'hd' in dep_metrics or 'dd' in dep_metrics else None
        results = {metric: [] for metric in dep_metrics}


        for word in sent:
            if word['head'] != 0:
                if 'dd' in dep_metrics:
                    results['dd'].append(self._dd(word))  
                if 'hd' in dep_metrics:
                    results['hd'].append(self._hd(word, word_index_by_id))
                if 'ddir' in dep_metrics:
                    results['ddir'].append(self._ddir(word)) 
                if 'v' in dep_metrics:
                    results['v'].append(self._v(word, sent)) 
                if 'l' in dep_metrics:
                    results['l'].append(self._l(word, sent))
            elif word['head'] == 0:
                if 'dd' in dep_metrics:
                    results['dd'].append(0)  
                if 'hd' in dep_metrics:
                    results['hd'].append(0)
                if 'ddir' in dep_metrics:
                    results['ddir'].append(0) 
                if 'v' in dep_metrics:
                    results['v'].append(self._v(word, sent)) 
                if 'l' in dep_metrics:
                    results['l'].append(self._l(word, sent))
            results['id'].append(word['id'])
            results['sent_id'].append(sent_id)
        return results
    
    def sent_info(self, sent: List[Dict[str, Any]], metrics: List[str] = None):
        if metrics is None:
            metrics = self.sent_metrics

        dep_metrics = list(set([self.projection[metric] for metric in metrics if metric in self.projection]))
        word_index_by_id = {word['id']: word for word in sent if type(word['id']) != tuple } if 'hd' in dep_metrics or 'dd' in dep_metrics else None
        temp_results = {metric: [] for metric in dep_metrics+['id','form','dpos','head','gform','gpos','deprel']}

        for word in sent:
            if word['head'] != 0:
                if 'dd' in dep_metrics:
                    temp_results['dd'].append(self._dd(word))  
                if 'hd' in dep_metrics:
                    temp_results['hd'].append(self._hd(word, word_index_by_id))
                if 'ddir' in dep_metrics:
                    temp_results['ddir'].append(self._ddir(word)) 
                if 'v' in dep_metrics:
                    temp_results['v'].append(self._v(word, sent)) 
                if 'l' in dep_metrics:
                    temp_results['l'].append(self._l(word, sent))
                temp_results['id'].append(word['id'])
                temp_results['form'].append(word['form'])
                temp_results['dpos'].append(word['upos'])
                temp_results['head'].append(word['head'])
                try:
                    temp_results['gform'].append(sent[word['head']-1]['form'])
                except:
                    temp_results['gform'].append('ERROR_ANNOTATION')
                try:
                    temp_results['gpos'].append(sent[word['head']-1]['upos'])
                except:
                    temp_results['gpos'].append('ERROR_ANNOTATION')
                temp_results['deprel'].append(word['deprel'])
            elif word['head'] == 0:
                if 'dd' in dep_metrics:
                    temp_results['dd'].append(0)  
                if 'hd' in dep_metrics:
                    temp_results['hd'].append(0)
                if 'ddir' in dep_metrics:
                    temp_results['ddir'].append(0) 
                if 'v' in dep_metrics:
                    temp_results['v'].append(self._v(word, sent)) 
                if 'l' in dep_metrics:
                    temp_results['l'].append(self._l(word, sent))
                temp_results['id'].append(word['id'])
                temp_results['form'].append(word['form'])
                temp_results['dpos'].append(word['upos'])
                temp_results['head'].append(word['head'])
                temp_results['gform'].append('ROOT')
                temp_results['gpos'].append('ROOT')
                temp_results['deprel'].append(word['deprel'])

        sent_data = {metric: [] for metric in metrics}
        metric_functions = {
            'mdd': lambda: sum([i for i in temp_results['dd'] if i > 0]) / max(1, len(temp_results['dd']) - 1),
            'mhd': lambda: sum([i for i in temp_results['hd'] if i > 0]) / max(1, len(temp_results['hd']) - 1),
            'mhdd': lambda: (len(temp_results['hd'])-1) / (max(temp_results['hd'])+1),
            'tdl': lambda: sum([i for i in temp_results['dd']]),
            'sl': lambda: len(temp_results['dd']),
            'vk': lambda: np.var(temp_results['v']),
            'tw': lambda: temp_results['hd'].count(max(temp_results['hd'], key=temp_results['hd'].count))+1,
            'th': lambda: max(temp_results['hd'])+1,
            'hi': lambda: temp_results['ddir'].count(1),
            'hf': lambda: temp_results['ddir'].count(-1),
            'ed_star': lambda: self.edit_distance_from_baseline(temp_results['head'],'star'),
            'ed_chain': lambda: self.edit_distance_from_baseline(temp_results['head'],'chain'),
            'ed_nest': lambda: self.edit_distance_from_baseline(temp_results['head'],'nest'),
            'sdd': lambda: sum(i for i in temp_results['dd']),
            'shd': lambda: sum(i for i in temp_results['hd']),
            'rsdd': lambda: None if temp_results['dd'].index(0) == len(temp_results['dd'])-1 else sum(temp_results['dd'][temp_results['dd'].index(0)+1:]),
            'lsdd': lambda: None if temp_results['dd'].index(0) == 0 else sum(temp_results['dd'][:temp_results['dd'].index(0)]),
            'lshd': lambda: None if temp_results['hd'].index(0) == len(temp_results['hd'])-1 else sum(temp_results['hd'][temp_results['hd'].index(0)+1:]),
            'rshd': lambda: None if temp_results['hd'].index(0) == 0 else sum(temp_results['hd'][:temp_results['hd'].index(0)]),
            'sldc': lambda: None if temp_results['ddir'].index(0) == 0 else temp_results['ddir'][:temp_results['ddir'].index(0)].count(-1),
            'srdc': lambda: None if temp_results['ddir'].index(0) == len(temp_results['ddir'])-1 else temp_results['ddir'][temp_results['ddir'].index(0)+1:].count(1),
            # 'vl': lambda: np.var(temp_results['l']),
            'lbl':lambda: temp_results['dd'].index(0),
            'rbl':lambda: len(temp_results['dd'])-temp_results['dd'].index(0)-1,
            'vl': lambda: self._vl(temp_results['head']),
            'rd': lambda: temp_results['dd'].index(0)+1,
            'nrd': lambda:(temp_results['dd'].index(0)+1)/len(temp_results['dd']),
            'ldc': lambda: None if temp_results['ddir'].index(0) == 0 else temp_results['ddir'][:temp_results['ddir'].index(0)].count(-1)/len(temp_results['ddir'][:temp_results['ddir'].index(0)]),
            'rdc': lambda: None if temp_results['ddir'].index(0) == len(temp_results['ddir'])-1 else temp_results['ddir'][temp_results['ddir'].index(0)+1:].count(1)/len(temp_results['ddir'][temp_results['ddir'].index(0)+1:]),
            'rmdd': lambda: None if temp_results['dd'].index(0) == len(temp_results['dd'])-1 else np.mean(temp_results['dd'][temp_results['dd'].index(0)+1:]),
            'lmdd': lambda: None if temp_results['dd'].index(0) == 0 else np.mean(temp_results['dd'][:temp_results['dd'].index(0)]),
            'lmhd': lambda: None if temp_results['hd'].index(0) == len(temp_results['hd'])-1 else np.mean(temp_results['hd'][temp_results['hd'].index(0)+1:]),
            'rmhd': lambda: None if temp_results['hd'].index(0) == 0 else np.mean(temp_results['hd'][:temp_results['hd'].index(0)]),
            'mndl': lambda: np.mean([temp_results['hd'][i - 1]for i in temp_results['id']if i not in temp_results['head']]),
            'ndd': lambda: abs(np.log((sum(i for i in temp_results['dd'] if i > 0) / max(1, len(temp_results['dd']) - 1)) / (np.sqrt((temp_results['dd'].index(0)+1) * max(1, len(temp_results['dd'])))))),
        }

        for metric in metrics:
            if metric in metric_functions:
                sent_data[metric] = metric_functions[metric]()
        return sent_data

    
    def calculate_sent_metrics(self, metrics: List[str] = None) -> Dict[str, List]:
        if metrics is None:
            metrics = self.sent_metrics
        
        dep_metrics = list(set([self.projection[metric] for metric in metrics if metric in self.projection]))
        dep_data = self.calculate_dep_metrics(dep_metrics)
        sent_data = {metric: [] for metric in metrics}

        metric_functions = {
            'mdd': lambda: [sum(i for i in j if i > 0) / max(1, len(j) - 1) for j in dep_data['dd']],
            'mhd': lambda: [sum(i for i in j if i > 0) / max(1, len(j) - 1) for j in dep_data['hd']],
            'mhdd': lambda: [(len(j)-1) / (max(j)+1) for j in dep_data['hd']],
            'tdl': lambda: [sum(j) for j in dep_data['dd']],
            'sl': lambda: [len(j) for j in dep_data['id']],
            'mv': lambda: [sum(j) / len(j) for j in dep_data['v']],
            'vk': lambda: [np.var(j) for j in dep_data['v']],
            'tw': lambda: [j.count(max(j, key=j.count))+1 for j in dep_data['hd']],
            'th': lambda: [max(j)+1 for j in dep_data['hd']],
            'hi': lambda: [j.count(1) for j in dep_data['ddir']],
            'hf': lambda: [j.count(-1) for j in dep_data['ddir']],
            # 'vl': lambda: [np.var(j) for j in dep_data['l']],
            'ed_star': lambda: [self.edit_distance_from_baseline(j,'star') for j in dep_data['head']],
            'ed_chain': lambda: [self.edit_distance_from_baseline(j,'chain') for j in dep_data['head']],
            'ed_nest': lambda: [self.edit_distance_from_baseline(j,'nest') for j in dep_data['head']],
            'vl': lambda: [self._vl(j) for j in dep_data['head']],
            'rd': lambda: [(j.index(0)+1) for j in dep_data['dd']],
            'vl': lambda: [self._vl(j) for j in dep_data['head']],
            'rd': lambda: [(j.index(0)+1) for j in dep_data['dd']],
            'lbl': lambda: [j.index(0) for j in dep_data['dd']],
            'rbl': lambda: [(len(j)-(j.index(0)+1)) for j in dep_data['dd']],
            'ldc': lambda: [None if j.index(0) == 0 else j[:j.index(0)].count(-1)/len(j[:j.index(0)]) for j in dep_data['ddir']],
            'rdc': lambda: [None if j.index(0) == len(j)-1 else j[j.index(0)+1:].count(1)/len(j[j.index(0)+1:]) for j in dep_data['ddir']],
            'lmdd': lambda:[None if j.index(0) == len(j)-1 else np.mean(j[:j.index(0)]) for j in dep_data['dd']],
            'rmdd': lambda:[None if j.index(0) == 0 else np.mean(j[j.index(0)+1:]) for j in dep_data['dd']],
            'rmhd': lambda:[None if j.index(0) == len(j)-1 else np.mean(j[j.index(0)+1:]) for j in dep_data['hd']],
            'lmhd': lambda:[None if j.index(0) == 0 else np.mean(j[:j.index(0)]) for j in dep_data['hd']],
            'mndl':lambda: [np.mean([d[2][i - 1]for i in d[0]if i not in d[1]]) for d in zip(dep_data['id'],dep_data['head'],dep_data['hd'])],
            'ndd': lambda: [abs(np.log((sum(i for i in j if i > 0) / max(1, len(j) - 1)) / (np.sqrt((j.index(0)+1) * max(1, len(j)))))) for j in dep_data['dd']],
        }

        for metric in metrics:
            if metric in metric_functions:
                sent_data[metric] = metric_functions[metric]()

        return sent_data
    
    def calculate_text_metrics(self, metrics: List[str] = None) -> Dict[str, List]:
        if metrics is None:
            metrics = self.text_metrics

        dep_metrics = list(set([self.projection[metric] for metric in metrics if metric in self.projection]))
        dep_data = self.calculate_dep_metrics(dep_metrics)
        
        text_data = {}

        def lbl(dds):
            root_idx = dds.index(0)+1
            return root_idx-1
        def rbl(dds):
            root_idx = dds.index(0)+1
            return len(dds)-root_idx
        def ldc(ddirs_list):
            left_count = 0
            count = 0
            for ddirs in ddirs_list:
                root_idx = ddirs.index(0)
                if root_idx > 0:
                    left_count += ddirs[:root_idx].count(-1)
                    count += len(ddirs[:root_idx])
            return left_count / count if count > 0 else None
        def rdc(ddirs_list):
            right_count = 0
            count = 0
            for ddirs in ddirs_list:
                root_idx = ddirs.index(0)
                if root_idx < len(ddirs)-1:
                    right_count += ddirs[root_idx+1:].count(1)
                    count += len(ddirs[root_idx+1:])
            return right_count / count if count > 0 else None
            
        def lmd(ds_list):
            total_sum = 0
            total_len = 0
            for ds in ds_list:
                root_idx = ds.index(0)
                if root_idx > 0:
                    seg = ds[:root_idx]
                    total_sum += sum(seg)
                    total_len += root_idx
            return total_sum / total_len if total_len > 0 else float('nan')
        
        def rmd(ds_list):
            total_sum = 0
            total_len = 0
            for ds in ds_list:
                root_idx = ds.index(0)
                if root_idx < len(ds)-1:
                    seg = ds[root_idx+1:]
                    total_sum += sum(seg)
                    total_len += root_idx
            return total_sum / total_len if total_len > 0 else float('nan')

        
        metric_calculators = {
            'mdd': lambda: sum(sum(dep_data['dd'], [])) / max(1,(len(sum(dep_data['dd'], []))-len(dep_data['dd']))),
            'mhd': lambda: sum(sum(dep_data['hd'], [])) / max(1,(len(sum(dep_data['hd'], []))-len(dep_data['hd']))),
            'mhdd': lambda: sum((len(j)-1) / (max(j)+1) for j in dep_data['hd']) / len(dep_data['hd']),
            'mtdl': lambda: sum(sum(j) for j in dep_data['dd']) / len(dep_data['dd']),
            'msl': lambda: sum(len(j) for j in dep_data['dd']) / len(dep_data['dd']),
            'mv': lambda: sum(sum(dep_data['v'], [])) / len(sum(dep_data['v'], [])),
            'vk': lambda: sum(np.var(j) for j in dep_data['v']) / len(dep_data['v']),
            'mtw': lambda: sum(j.count(max(j, key=j.count))+1 for j in dep_data['hd']) / len(dep_data['hd']),
            'mth': lambda: sum(max(j)+1 for j in dep_data['hd']) / len(dep_data['hd']),
            'hi': lambda: sum(j.count(1) for j in dep_data['ddir']) / max(1,(len(sum(dep_data['ddir'], []))-len(dep_data['ddir']))),
            'hf': lambda: sum(j.count(-1) for j in dep_data['ddir']) / max(1,(len(sum(dep_data['ddir'], []))-len(dep_data['ddir']))),
            'vl':lambda: sum(np.var(j) for j in dep_data['l']) / len(dep_data['l']),
            'lbl': lambda: sum(lbl(j) for j in dep_data['dd']) / len(dep_data['dd']),
            'rbl': lambda: sum(rbl(j) for j in dep_data['dd']) / len(dep_data['dd']),
            'ldc': lambda: ldc(dep_data['ddir']),
            'rdc': lambda: rdc(dep_data['ddir']),
            'lmdd': lambda: lmd(dep_data['dd']),
            'rmdd': lambda: rmd(dep_data['dd']),
            'lmhd': lambda: lmd(dep_data['hd']),
            'rmhd': lambda: rmd(dep_data['hd']),
            # 'lmdd': lambda: sum([lmdd(j) for j in dep_data['dd'] if lmdd(j) != None]) / len(dep_data['dd']),
            # 'rmdd': lambda: sum([rmdd(j) for j in dep_data['dd'] if rmdd(j) != None]) / len(dep_data['dd']),
            # 'lmhd': lambda: sum([lmhd(j) for j in dep_data['hd'] if lmhd(j) != None]) / len(dep_data['hd']),
            # 'rmhd': lambda: sum([rmhd(j) for j in dep_data['hd'] if rmhd(j) != None]) / len(dep_data['hd']),
            'mrd': lambda: sum(j.index(0)+1 for j in dep_data['dd']) / len(dep_data['dd']),
            'nrd': lambda: sum((j.index(0)+1)/len(j) for j in dep_data['dd']) / len(dep_data['dd']),
            'mndl':lambda: np.mean([np.mean([d[2][i - 1]for i in d[0]if i not in d[1]]) for d in zip(dep_data['id'],dep_data['head'],dep_data['hd'])]),
            'ndd': lambda: sum(abs(np.log((sum(i for i in j if i > 0) / max(1, len(j) - 1)) / (np.sqrt((j.index(0)+1) * max(1, len(j)))))) for j in dep_data['dd']) / len(dep_data['dd']),
        }
    
        for metric in metrics:
            if metric in metric_calculators:
                text_data[metric] = metric_calculators[metric]()

        return text_data


    def calculate_distributions(self, metrics: List[str] = None, normalize: bool = False) -> Dict[str, Tuple[List, List]]:
        if metrics is None:
            metrics = self.distribution_metrics

        dep_metrics = list(set(self.projection.get(metric, metric) for metric in metrics))
        dep_data = self.calculate_dep_metrics(dep_metrics)

        distributions = {}
        metric_functions = {
            'sl': lambda: [len(j) for j in dep_data['id']],
            'tw': lambda: [j.count(max(j, key=j.count))+1 for j in dep_data['hd']],
            'th': lambda: [max(j)+1 for j in dep_data['hd']],
            'rd': lambda: [j.index(0) for j in dep_data['head']]
        }

        def normalize_data(data):
            return [i[1] / sum(j[1] for j in data) for i in data] if normalize else [i[1] for i in data]

        for metric in metrics:
            if metric not in ['deprel', 'pos']:
                data = sorted(Counter(sum(dep_data[metric], [])).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)

            elif metric in ['deprel', 'pos']:
                data = sorted(Counter(sum(dep_data[self.projection.get(metric, metric)], [])).items(),key=lambda x: x[1], reverse=True)
                y = normalize_data(data)
                distributions[metric] = (range(1, len(y) + 1), y)

            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)

        return distributions


    def calculate_pvp(self, input: Optional[str] = None, target:str='deprel',normalize:bool=True):

        dependents = defaultdict(int)
        governors = defaultdict(int)

        def process_word(word, sentence):
            if target == 'deprel':
                dependent = [w['deprel'] for w in sentence if w['head'] == word['id'] and w['deprel'] not in stopdeps]
                for dep in dependent:
                    dependents[dep] += 1
                governors[word['deprel']] += 1

            if target == 'pos':
                dependent = [w['upos'] for w in sentence if w['head'] == word['id'] and w['deprel'] not in stopdeps]
                for dep in dependent:
                    dependents[dep] += 1
                governor = [w['upos'] for w in sentence if w['id'] == word['head'] and w['deprel'] not in stopdeps]
                for gov in governor:
                    governors[gov] += 1

        for sentence in self.treebank:
            for word in sentence:
                if word['deprel'] not in stopdeps and (input is None or word['upos'] == input):
                    process_word(word, sentence)

        total_deps = sum(dependents.values())
        total_govs = sum(governors.values())
        if normalize:
            govs = sorted(((k, v / total_govs) for k, v in governors.items()), key=lambda x: x[1], reverse=True)
            deps = sorted(((k, v / total_deps) 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': deps, 'act as a dep': govs}
    
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):
    treebank = DepValAnalyzer(treebank)
    text_data = treebank.calculate_text_metrics()
            
    dep_data = treebank.calculate_dep_metrics()
    dep_data = {x:sum(y,[]) for x,y in dep_data.items()}
    dep_data = pd.DataFrame(dep_data)

    sent_data = treebank.calculate_sent_metrics()
    sent_data = pd.DataFrame(sent_data)  

    pvp_data = treebank.calculate_pvp()
    pvp_data = _pvp2df(pvp_data)  

    distribution_dict = {}
    distributions = treebank.calculate_distributions(normalize=normalize)
    for m in ['dd','hd','sl','v','tw','th','rd','deprel','pos']:
        data = pd.DataFrame(distributions[m]).T
        distribution_dict[m] = data
    
    return text_data,dep_data,sent_data,pvp_data,distribution_dict

def analyze(treebank_path:str,out_path:str,normalize:bool=True):

    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')]
        text_csv = pd.DataFrame(columns=['mdd','mhd','mhdd','mtdl','msl','mv','vk','mtw','mth','hi','hf'])
        file_names = [os.path.basename(t).split('.')[0] for t in treebanks]
        files = [os.path.join(out_path, f) for f in file_names]
        for i,t in enumerate(treebanks):      
            os.makedirs(files[i],exist_ok=True)     
            treebank = open(t, encoding='utf-8')
            text_metrics,dep_metrics,sent_metrics,pvp_metrics,distribution_dict = getDepValFeatures(treebank)
                    
            text_csv.loc[i,:] = text_metrics   
            dep_metrics.round(2).to_csv(os.path.join(files[i],f'dep_metrics.csv'),index=False)   
            sent_metrics.round(2).to_csv(os.path.join(files[i],f'sent_metrics.csv'),index=False)   
            pvp_metrics.round(4).to_csv(os.path.join(files[i],f'pvp.csv'),index=False)    

            for m in distribution_dict:
                distribution_dict[m].to_csv(os.path.join(files[i],f'{m}_distribution.csv'),index=False,header=False) 
        
        text_csv = text_csv.astype(float).round(2)
        text_csv['treebank'] = file_names
        text_csv.to_csv(os.path.join(out_path, f'text_metrics.csv'), index=False)

class RandomLanguage:

    def __init__(self):
        self.params = ['sl','rd']
        self.metrics = ['mdd','mhd','tw','th','vk','vl']
        fake_tree = [{'id':1,'form':'token','lemma':'_','upos':'_','xpos':'_','feats':'_',
                        'head':0,'deprel':'root','deps':'_','misc':'_'}]
        self.dvanalyzer = DepValAnalyzer(fake_tree, cleaned=True)
    
    # def DTGA(self,sentlength, cross=False, root=None):
    #     """
    #     Dependency Tree Generator Algorithm (DTGA_v3)
    #     生成随机依存结构：
    #         - cross=False → 无交叉projective树
    #         - cross=True  → 允许交叉但保持连通的随机树
    #     """

    #     node_sequence = list(range(1, sentlength + 1))
        
    #     # Step 1. 选根节点
    #     if root is None:
    #         root = random.choice(node_sequence)
        
    #     links = []

    #     # --- cross=False: 生成严格projective依存树 ---
    #     if not cross:
    #         def build_subtree(start, end, root_idx):
    #             if start > end:
    #                 return
    #             left_nodes = list(range(start, root_idx))
    #             right_nodes = list(range(root_idx + 1, end + 1))
    #             if left_nodes:
    #                 left_head = random.choice(left_nodes)
    #                 links.append((left_head, root_idx))
    #                 build_subtree(start, root_idx - 1, left_head)
    #             if right_nodes:
    #                 right_head = random.choice(right_nodes)
    #                 links.append((right_head, root_idx))
    #                 build_subtree(root_idx + 1, end, right_head)

    #         build_subtree(1, sentlength, root)
    #         links.append((root, 0))
    #         return sorted(links)

    #     # --- cross=True: 生成连通但可交叉的树 ---
    #     nodes_in_tree = [root]
    #     remaining_nodes = [i for i in node_sequence if i != root]

    #     while remaining_nodes:
    #         dep = random.choice(remaining_nodes)
    #         head_candidate = random.choice(nodes_in_tree)
    #         links.append((dep, head_candidate))
    #         nodes_in_tree.append(dep)
    #         remaining_nodes.remove(dep)

    #     links.append((root, 0))
    #     return sorted(links)
    # def DTGA(self,sentlength,cross=False,head=None): 
    #     # 创建一个节点序列列表，包含从1到句子长度(sentlength)的所有整数
    #     node_sequence = list(range(1,sentlength+1)) 
    #     # 检查头节点(head)是否为空
    #     if head is None: 
    #         head = random.choice(node_sequence) 
    #     if cross == True: 
    #         links = {i:random.choice([j for j in node_sequence if j!= i]) for i in node_sequence if i!= head} 
    #         links = [(k,v) for k,v in links.items()]
    #     elif cross == False: 
    #         links = [] 
    #         nodes_in_tree = [head] 
    #         node_not_in_tree = [i for i in node_sequence if i not in nodes_in_tree] 
    #         while node_not_in_tree: 
    #             random_node = random.choice(node_not_in_tree) 
    #             if not links: 
    #                 links.append((random_node,head)) 
    #             else: 
    #                 potential_governor = set(nodes_in_tree) 
    #                 for l in links: 
    #                     if random_node in range(min(l[0],l[1]),max(l[0],l[1])+1): 
    #                         potential_governor.intersection_update({i for i in nodes_in_tree if i in range(min(l[0],l[1]),max(l[0],l[1])+1)}) 
    #                     else: 
    #                         potential_governor.intersection_update({i for i in nodes_in_tree if i not in range(min(l[0],l[1])+1,max(l[0],l[1]))}) 
    #                 links.append((random_node,random.choice(list(potential_governor)))) 
    #             nodes_in_tree.append(random_node) 
    #             node_not_in_tree.remove(random_node) 
    #         links.append((head,0)) 
    #     return links

    def DTGA(self,sl,cross=False,root=None):

        def rl1(sl,root=None):
            V = list(range(1,sl+1))
            DG_r = nx.DiGraph()
            if sl == 1:
                DG_r.add_node(V[0])
                root = V[0]
                return [(root,0)]
            Vtemp = V[:]
            tmp = V[:]
            if root == None:
                root = random.choice(V)
            Vtemp.remove(root)
            while(len(Vtemp)!=0):
                j = random.choice(Vtemp)
                tmp.remove(j)
                if DG_r.has_node(j):
                    sucs_dict = nx.dfs_successors(DG_r,j)
                    if sucs_dict != {}:
                        for key,value in sucs_dict.items():
                            for val in value:
                                tmp.remove(val)
                    else:
                        continue
                i = random.choice(tmp)
                DG_r.add_edge(i,j)
                Vtemp.remove(j)
                tmp = V[:]
            edges = sorted([(e[1],e[0]) for e in DG_r.edges()] + [(root,0)])
            return edges
        def rl2(sl,root=None):
            V = list(range(1,sl+1))
            DG_r = nx.DiGraph()
            if sl == 1:
                DG_r.add_node(V[0])
                root = V[0]
                return [(root,0)]
            Vtemp = V[:]
            tmp = V[:]
            if root == None:
                root = random.choice(V)
            Vtemp.remove(root)
            while(len(Vtemp)!=0):
                j = random.choice(Vtemp)
                segment = []
                if j>root:               #1.确定j(dependent)在root的哪边
                    tmp = V[V.index(root):]
                if j<root:
                    tmp = V[:V.index(root)+1]
                if DG_r.has_node(j):      #2.已生成的图中是否有j，即j为governers(j->nodes),从而排除这些节点作为j的governers的可能
                    sucs_dict = nx.dfs_successors(DG_r,j)
                    if sucs_dict != {}:
                        for key,value in sucs_dict.items():
                            for val in value:
                                segment.append(val)  
                        segment.append(j)
                        maxn = max(segment)
                        minn = min(segment)
                        for val_segment in range(minn,maxn+1):#去除j与其后继结点。governers候选者需要删除j以及后继结点构成的所有的节点。不仅去掉了回环的可能，也去掉了交叉的可能。
                            if val_segment in tmp:
                                tmp.remove(val_segment)           #这是在j与其后继结点形成的段内。
                if DG_r.edges() != []:#3.去除与j有间隔的弧，其可能和j发生交叉的节点
                    minarclen = len(V)-1
                    inarc = ()
                    for edge in DG_r.edges():  #发现包含j的区间的最短弧长
                        tmp1 = []
                        if min(edge[0],edge[1])<j<max(edge[0],edge[1]):
                            arclen = abs(edge[0]-edge[1])
                            if arclen < minarclen:
                                minarclen = arclen
                                tmp1 = V[V.index(min(edge[0],edge[1])):(V.index(max(edge[0],edge[1]))+1)]
                                tmp = list(set(tmp).intersection(set(tmp1)))
                                inarc = edge
                    for edgeinterval in DG_r.edges():    # j在某些弧的外部，产生I型交叉的情况
                        if min(edgeinterval[1],edgeinterval[0])>j:
                            for ee in range((min(edgeinterval[1],edgeinterval[0])+1),max(edgeinterval[1],edgeinterval[0])):
                                if ee in tmp:
                                    tmp.remove(ee)
                        if max(edgeinterval[1],edgeinterval[0])<j:
                            for ee in range((min(edgeinterval[1],edgeinterval[0])+1),max(edgeinterval[1],edgeinterval[0])):
                                if ee in tmp:
                                    tmp.remove(ee)
                        if edgeinterval[1]<edgeinterval[0]<j:    # j在某些弧的外部，产生II型交叉的情况
                            if edgeinterval[1] in tmp:
                                tmp.remove(edgeinterval[1])
                        if edgeinterval[1]>edgeinterval[0]>j:
                            if edgeinterval[1] in tmp:
                                tmp.remove(edgeinterval[1])
                if j in tmp:
                    tmp.remove(j)
                i = random.choice(tmp)
                DG_r.add_edge(i,j)
                Vtemp.remove(j)
                tmp = V[:]
            edges = sorted([(e[1],e[0]) for e in DG_r.edges()] + [(root,0)])
            return edges
        if cross:
            edges = rl1(sl,root=root)
        elif not cross:
            edges = rl2(sl,root=root)
        return edges 

    def _sl(self,sent):
        return len(sent)
    
    def _rd(self,sent):
        return [w['id'] for w in sent if w['head'] == 0][0]

    def links_to_conllu(self,links):
        sent = []
        for l in links:
            token = {}
            token['form'] = 'token'
            token['lemma'] = '_'
            token['upos'] = '_'
            token['xpos'] = '_'
            token['feats'] = '_'
            token['deps'] = '_'
            token['misc'] = '_'
            token['id'] = l[0]
            token['head'] = l[1]
            if l[1] != 0:
                token['deprel'] = 'non-root'
            else:
                token['deprel'] = 'root'
            sent.append(token)
        return sent 

    def generate_random_trees(self,sl=10,root=None,cross=False):
        
        links = self.DTGA(sl,cross=cross,root=root)
        sent = self.links_to_conllu(links)
        return sent
    
    def generate_chains(self,sl=10,root=None,cross=False):
        if root is None:
            root = random.choice(range(1,sl+1))
        left_links = [(i,i+1) for i in range(1,root)]
        right_links = [(i,i-1) for i in range(root+1,sl+1)]
        links = left_links +[(root,0)]+ right_links
        sent = self.links_to_conllu(links)
        return sent
    
    def generate_stars(self,sl=10,root=None,cross=False):
        if root is None:
            root = random.choice(range(1,sl+1))
        links = [(i,0) if i == root else (i,root) for i in range(1, sl + 1)]
        sent = self.links_to_conllu(links)
        return sent
    
    def generate_nests(self,sl=10,root=None,cross=False):
        if root is None:
            root = random.choice([1, sl])
            
        links = []
        l, r = (1, sl) if root == 1 else (sl, 1)
        step = 1 if root == 1 else -1

        for i in range(sl - 1):
            links.append((l, r))
            l, r = r, l + step if i % 2 == 0 else l - step
        links = links if root == 1 else [(i[1],i[0]) for i in links]
        if root == 1:links = [(i[1],i[0]) for i in links]
        links.append((root, 0))
        links = sorted(links)
        sent = [i[1] for i in links]
        sent = self.links_to_conllu(links)
        return sent

    def sampling(self,sl_list,roots=None,cross=False,sample_size=100,pool=False):
        if roots is None:
            roots = [None]*len(sl_list)
        paras = zip(sl_list,roots)
        sents = []

        if pool is False:
            sents = []
            for slr in paras:
                for i in range(sample_size):
                    sents.append(self.generate_random_trees(sl=slr[0],root=slr[1],cross=cross))
        elif pool is True:
            num_processes = mp.cpu_count() 
            with mp.Pool(processes=num_processes) as pool:
                args = [(slr[0],slr[1],cross) for slr in paras for _ in range(sample_size)]
                sents = pool.starmap(self.generate_random_trees, args)
        elif type(pool) is int:
            num_processes = pool
            with mp.Pool(processes=num_processes) as pool:
                args = [(slr[0],slr[1],cross) for slr in paras for _ in range(sample_size)]
                sents = pool.starmap(self.generate_random_trees, args)
        return sents
    
    def _sample_sent(self,sl, root, cross, metrics,star=False,chain=False,nest=False):
            analyzer = self.dvanalyzer
            if chain:
                sent = self.generate_chains(sl=sl, root=root,cross=cross)
            elif star:
                sent = self.generate_stars(sl=sl, root=root,cross=cross)
            elif nest:
                sent = self.generate_nests(sl=sl, root=root, cross=cross)
            else:
                sent = self.generate_random_trees(sl=sl, root=root, cross=cross)
            
            sent_metrics = analyzer.sent_info(sent, metrics=metrics)   
            return sent_metrics
    
    def _sample_many(self, sl, root, cross, metrics,sample_size, star=False, chain=False, nest=False):
        
        analyzer = self.dvanalyzer

        results = []

        for i in range(sample_size):
            if chain:
                sent = self.generate_chains(sl=sl, root=root, cross=cross)
            elif star:
                sent = self.generate_stars(sl=sl, root=root, cross=cross)
            elif nest:
                sent = self.generate_nests(sl=sl, root=root, cross=cross)
            else:
                sent = self.generate_random_trees(sl=sl, root=root, cross=cross)

            sent_metrics = analyzer.dep_info(sent, metrics=metrics)

            results.append(sent_metrics)

        return results
    
    def sampling_sent_metrics(self, sl_list, roots=None, cross=False, sample_size=100, metrics=None, pool=False,star=False,chain=False,nest=False):
        result = {m: [] for m in metrics}

        all_results = []

        if not pool:
            if roots is not None:
                for slr in zip(sl_list,roots):
                    for i in range(sample_size):
                        all_results.append(self._sample_sent(slr[0], slr[1], cross, metrics))
            elif roots is None:
                for sl in sl_list:
                    for i in range(sample_size):
                        all_results.append(self._sample_sent(sl, None, cross, metrics))
        else:
            num_processes = mp.cpu_count() if pool is True else pool
            with mp.Pool(processes=num_processes) as p:
                if roots is not None:
                    slr = zip(sl_list,roots)
                    args = [(i[0], i[1], cross, metrics, star, chain,nest) for i in slr for _ in range(sample_size)]
                elif roots is None:
                    args = [(sl, roots, cross, metrics, star, chain, nest) for sl in sl_list for _ in range(sample_size)]
                all_results = p.starmap(self._sample_sent, args)

        for res in all_results:            
            for m in metrics:
                result[m].append(res[m])
        return result

    def _sample_dep(self,sl, root, cross, metrics,sent_id=0, star=False,chain=False,nest=False):
            analyzer = self.dvanalyzer
            if chain:
                sent = self.generate_chains(sl=sl, root=root,cross=cross)
            elif star:
                sent = self.generate_stars(sl=sl, root=root,cross=cross)
            elif nest:
                sent = self.generate_nests(sl=sl, root=root, cross=cross)
            else:
                sent = self.generate_random_trees(sl=sl, root=root, cross=cross)
            dep_metrics = analyzer.dep_info(sent, metrics=metrics,sent_id=sent_id)  
            dep_metrics['para'] = [(int(sl),None) if root==None else (int(sl),int(root)) for k in range(sl)]
            
            return dep_metrics
    
    def sampling_dep_metrics(self, sl_list, roots=None, cross=False, sample_size=100, metrics=None, pool=False,star=False,chain=False,nest=False):
        
        result = {m: [] for m in metrics+['para']}
        sent_id = list(range(len(sl_list)))

        all_results = []

        if not pool:
            if roots is not None:
                for slr in zip(sl_list,roots):
                    for i in range(sample_size):
                        all_results.append(self._sample_dep(slr[0], slr[1], cross, metrics, sent_id=i))
            elif roots is None:
                for sl in sl_list:
                    for i in range(sample_size):
                        all_results.append(self._sample_dep(sl, None, cross, metrics, sent_id=i))
        else:
            num_processes = mp.cpu_count() if pool is True else pool
            with mp.Pool(processes=num_processes) as p:
                if roots is not None:
                    # slr = list(zip(sl_list,roots))
                    args = [(sl_list[i], roots[i], cross, metrics, j*len(sl_list)+i, star, chain,nest) for i in sent_id for j in range(sample_size)]
                elif roots is None:
                    args = [(sl_list[i], roots, cross, metrics, j*len(sl_list)+i, star, chain, nest) for i in sent_id for j in range(sample_size)]
                all_results = p.starmap(self._sample_dep, args)

        for deps in all_results:    
            for m in metrics+['para']:
                result[m] += deps[m]
        return result    

    def _get_metric(self,analyzer,sent,metrics=None):
        if metrics is None:
            metrics = self.metrics
        else:
            sent_metrics = analyzer.sent_info(sent,metrics=metrics)
            return sent_metrics
    def _get_metrics(self,sents,metrics=None):
        if metrics is None:
            metrics = self.metrics
        else:
            analyzer = DepValAnalyzer(sents,cleaned=True)
            sent_metrics = analyzer.calculate_sent_metrics(metrics=metrics)
        return sent_metrics
    
    def get_metrics(self,sents,metrics=None,pool=False):
        if metrics is None:
            metrics = self.metrics
        if pool is False:
            sent_metrics = self._get_metrics(sents,metrics=metrics)
        elif pool is True:
            num_processes = mp.cpu_count()
            sent_metrics = {metric: [] for metric in metrics}
             
            fake_tree = [{'id':1,'form':'token','lemma':'_','upos':'_','xpos':'_','feats':'_','head':0,'deprel':'root','deps':'_','misc':'_'}]
            analyzer = DepValAnalyzer(fake_tree,cleaned=True)
            with mp.Pool(processes=num_processes) as pool:
                args = [(analyzer,sent,metrics) for sent in sents]
                for result in pool.starmap(self._get_metric, args):
                    for metric in metrics:
                        sent_metrics[metric].append(result[metric])
        elif type(pool) is int:
            num_processes = pool
            sent_metrics = {metric: [] for metric in metrics}
             
            fake_tree = [{'id':1,'form':'token','lemma':'_','upos':'_','xpos':'_','feats':'_','head':0,'deprel':'root','deps':'_','misc':'_'}]
            analyzer = DepValAnalyzer(fake_tree,cleaned=True)
            with mp.Pool(processes=num_processes) as pool:
                args = [(analyzer,sent,metrics) for sent in sents]
                for result in pool.starmap(self._get_metric, args):
                    for metric in metrics:
                        sent_metrics[metric].append(result[metric])
        return sent_metrics

    def map_ids(self,natural_sent_metrics:dict,random_sent_metrics,sample_size:int,label='random'):
        natural_data = pd.DataFrame(natural_sent_metrics)
        natural_data['langtype'] = 'natural'
        natural_data['ids'] = list(range(len(natural_sent_metrics['sl'])))
        def get_random_ids(rand, sample_size,label='random'):
            random_data = pd.DataFrame(rand)
            random_data['langtype'] = label
            rand_ids = []
            for i in range(len(natural_sent_metrics['sl'])):
                rand_ids += [i] * sample_size
            random_data['ids'] = rand_ids
            return random_data

        if type(random_sent_metrics) is dict:
            random_data = get_random_ids(random_sent_metrics, sample_size)
            mapped_sent_metrics = pd.concat([natural_data,random_data],ignore_index=True)

        elif type(random_sent_metrics) is list:
            mapped_sent_metrics = natural_data.copy()
            for i,r in enumerate(random_sent_metrics):
                random_data = get_random_ids(r, sample_size,label=f'{label[i]}')
                mapped_sent_metrics = pd.concat([mapped_sent_metrics,random_data],ignore_index=True)
        return mapped_sent_metrics


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': int(columns[3][i]),
                        'deprel': columns[2][i],
                        'deps': '_',
                        'misc': '_'
                    }
                    for i in range(len(columns[0]))
                ]
                
                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):
                treebank = open(t, encoding='utf-8')
                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]}.txt'))
    else:
        raise ValueError(f"The input path {treebank_path} is not a valid directory. It should be a directory containing treebank files.")





