#!/usr/bin/env python3

import re
import sys
import math
from collections import defaultdict, Counter
import argparse

from POSTagger import Tagger  # Wortart-Tagger
from Transducer import Transducer
from betonung import Betonung
from zeichenklassen import v, c, p, f, n
from silbifizierung import syllabify

# reduction_threshold = 500

vers_agreement = {
    'default'  :'',
    'reimpaare':'Kadenz+Hebigkeit',
    'strophen' :'Kadenz+Hebigkeit',
    'frei'     :'Kadenz+Hebigkeit',
}

def parse_arguments():
    parser = argparse.ArgumentParser(description='Versanalyse')
    parser.add_argument('eingabedatei',
                        help='Eingabedatei')
    parser.add_argument('-k', '--kostendatei',
                        default='lib/kosten.txt',
                        help='Kostendatei')
    parser.add_argument('-t', '--typ',
                        choices=vers_agreement.keys(),
                        default=next(iter(vers_agreement.keys())),
                        help='Verstyp')
    parser.add_argument('-e', '--elisionen',
                        action='store_true', default=False,
                        help='erlaubt mehr Elisionen')
    parser.add_argument('-f', '--formel', help='Vorgabe der Strophenformel')
    parser.add_argument('-v', '--verbose',
                        action='store_true', default=False,
                        help='gibt mehr Informationen aus')
    parser.add_argument('-n', '--keine-namen',
                        action='store_true', default=False,
                        help='gibt mehr Informationen aus')
    return parser.parse_args()


class Versanalyse:

    def __init__(self, args):
        print('Initialisierung ... ', end='', file=sys.stderr, flush=True)
        self.verbose = args.verbose
        self.strophenformel = [[t] for t in args.formel.split()] if args.formel else args.formel # Liste von Listen
        self.betonung = Betonung(not args.keine_namen, args.elisionen)
        self.normalisierung = self.lese_normalisierungen('lib/normalisierungen.txt') if args.elisionen else {}
        self.tagger = Tagger('lib/Tagger/parfile')
        self.transducer = Transducer('lib/transducer.txt', args.kostendatei)
        print('fertig\nAnalyse', file=sys.stderr)

        sys.stdout.reconfigure(encoding='utf-8')

        # Interpunktionszeichen, die der Tokenizer von den Wörtern abtrennt
        self.punctuation = '.,;:?!"`()„“»«–— \\-'

        
    @staticmethod
    def lese_normalisierungen(dateiname):
        normalisierung = {}
        with open(dateiname, encoding='utf-8') as file:
            for line in file:
                # Kommentare entfernen
                line = re.sub(r'\s*#.*', '', line.strip())
                if line:
                    try:
                        word, norm = line.split('\t')
                    except:
                        print(f'Fehler in Normalisierungsdatei: {line}', file=sys.stderr)
                    else:
                        normalisierung[word] = norm
        return normalisierung

        
    def normalisiere(self, line):
        ''' normalisiert einzelne Wörter wie "und" "umb" '''
        if not self.normalisierung:
            return line
        
        # Präfix-Klitika abtrennen wie bei z'alle
        line = re.sub(r"\b([sz]')(\w)", r'\1 \2', line)
        # Synkopen wie g'ahtet
        line = re.sub(rf"\b([g])'({v})", r'\1e\2', line)
        
        # Suffix-Klitika abtrennen wie bei mich's
        line = re.sub(r"(\w)('[szn])\b", r'\1 \2', line)
        
        # splitten
        tokens = re.findall(rf'[^{self.punctuation} ]+|[{self.punctuation} ]', line)

        # normalisieren
        new_tokens = [self.normalisierung.get(token, token)
                      for token in tokens]
    
        return ''.join(new_tokens)


    def tokenisiere(self, line):
        ''' zerlegt einen Satz in Tokens '''
        # Normalisierung von Zeichen
        line = line.strip().translate(str.maketrans('ẹ„“»«›‹–— ', 'e""""""-- '))

        # ' durch " ersetzen außer es ist Teil klitischer Ausdrücke
        line = re.sub(r"(\w)'(\w)", r'\1@\2', line).replace("'", '"').replace('@', "'")
        
        # klitische Ausdrücke:  z' ime --> z'ime
        line = re.sub(rf"\b({c}+') (\w)", r"\1\2", line)
        line = re.sub(rf"(\w) ('{c}+)\b", r"\1\2", line)
    
        # Satz in Tokens aufspalten und Interpunktion abtrennen
        tokens = re.findall(rf'[^{self.punctuation} ]+|[{self.punctuation}]', line)

        return tokens

    @staticmethod
    def generaltyp(typ):
        kadenz, hebigkeit = typ
        # ("3k","4") --> "4k"
        return hebigkeit + kadenz[-1]


    def align(self, result, zeile):
        # Normalisierungen anwenden
        zeile = self.normalisiere(zeile)

        # prevent index errors
        result += '\n'; zeile += '\n'
        
        new_result = ''
        flag = False
        i, k = 0, 0
        while result[i] != '\n' or zeile[k] != '\n':
            r, z,  = result[i], zeile[k]
            if r == z.lower():  # unverändertes Zeichen
                new_result += z
                i += 1; k += 1
            elif result[i:i+2] == 'ss' and z == 'ß':  # beßer --> besser
                new_result += 'ß'
                i += 2; k += 1
            elif not z.isalpha() and z != '\n':  # Interpunktion u.Ä.
                new_result += z
                k += 1
            elif r == '°':  # Betonungszeichen
                new_result += "'"
                i += 1
            elif r == '~':  # schwebende Betonung
                new_result += r
                i += 1
            elif z in 'eẹ' and r in 'eẹ':  # Elision
                new_result += r
                i += 1; k += 1
            elif r == ' ':  # fehlendes Leerzeichen in der Eingabezeile
                new_result += ' '
                i += 1
            else:
                sys.exit(f'Fehler in der Funktion align: {r}  {z}')
        return new_result

    
    def get_analysis(self, zeile, alle_silben, transducer_output):
        # join list elements and split on whitespace
        transducer_output = ''.join(transducer_output).split()

        result = []
        for silben, output in zip(alle_silben, transducer_output):
            wort = []
            for silbe, symbol in zip(silben, output):
                match symbol:
                    case 'H':
                        wort.append('°' + silbe)
                    case 'h':  # schwebende Betonung
                        wort.append('~' + silbe)
                    case '.':
                        wort.append(silbe.replace('e', 'ẹ'))
                    case _:
                        wort.append(silbe)
            result.append(wort)
        result = ' '.join(''.join(wort) for wort in result)
        return self.align(result, zeile)


    def print_analysis(self, versdict, verstyp):
        # verstyp: bspw. ('2k', '3')

        # Analysen mit minimalen Kosten bestimmen
        best_tuples = min(versdict['analysen'].values(), key=lambda x: x[0][0])
        min_kosten = best_tuples[0][0]

        # Analysen für den übergebenen Verstyp nachschlagen
        # und markieren, falls ihre Kosten nicht minimal sind
        if verstyp in versdict['analysen'].keys():
            tuples = versdict['analysen'][verstyp]
            optimal = '*' if tuples[0][0] > min_kosten and self.verbose else ''
        else:
            tuples, optimal = best_tuples, '**' if self.verbose else ''

        # Kadenz und Hebigkeit aus den Markern extrahieren
        _, output, markers, _ = tuples[0]
        kadenz, hebigkeit = markers[-2:]
        match = re.match(r'<([12]m|w|[23]k)><([1-9])h>$', kadenz+hebigkeit)
        if not match:
            sys.exit('Fehler: Kadenz und/oder Hebigkeit kann nicht bestimmt werden')
        kadenz, hebigkeit = match.group(1,2)
        code = self.generaltyp((kadenz, hebigkeit))
        if output[0] == 's':  # Senkung
            code = '.' + code

        # Generate the output
        if self.verbose:
            print(versdict['zeile'])
            print(*versdict['wörter'])
            print(*versdict['tags'])
            print('  '.join(' '.join(s) for s in versdict['silben']))
            print(versdict['analyseinput'])
        else:
            print(versdict['zeile'], end='\t')

        if not versdict['analysen']:
            print('keine erfolgreiche Analyse')
        else:
            for loss, output, markers, raw_output in tuples:
                ana = self.get_analysis(versdict['norm_zeile'], versdict['silben'], output)
                if self.verbose:
                    print(ana)
                    print(''.join(output))
                    print(code)
                    print(*raw_output, sep='')
                    print(*markers)
                    print('Kosten:', round(loss, 2))
                else:
                    print(ana, ''.join(output), code, sep='\t')
        if self.verbose:
            print()

            
    def print_all_analyses(self, strophen):
        for strophe in strophen:
            for versdict in strophe:
                self.print_analysis(versdict, versdict['verstyp'])
            print()
            

    def default_analyse(self, strophen):
        ''' keine Reimpaare und keine Strophen '''
        for strophe in strophen:
            for versdict in strophe:
                versdict['verstyp'] = ''


    def reimpaar_analyse(self, strophen):
        ''' analysierte einen Text bestehend aus Reimpaaren '''

        # Fehlerhafte Eingaben erkennen
        if any(len(strophe) % 2 == 1 for strophe in strophen):
            print('Fehler: Nicht alle Strophen haben gerade Länge!', file=sys.stderr)
            return

        for strophe in strophen:
            for i in range(0, len(strophe), 2):
                verse = strophe[i:i+2]  # nächstes Reimpaar
                ana1 = verse[0]['analysen']
                ana2 = verse[1]['analysen']
            
                # Iteration über die gemeinsamen Analysen
                analysen = set(ana1.keys()) & set(ana2.keys())
                if analysen:
                    # Die erste Verszeile erhält doppeltes Gewicht
                    total_loss = {verstyp: 2 * ana1[verstyp][0][0] + ana2[verstyp][0][0]
                                  for verstyp in analysen}
                    verstyp1 = verstyp2 = min(total_loss, key=total_loss.get)
                else:
                    loss1 = {verstyp: info[0][0] for verstyp, info in ana1.items()}
                    loss2 = {verstyp: info[0][0] for verstyp, info in ana2.items()}
                    verstyp1 = min(loss1, key=loss1.get)
                    verstyp2 = min(loss2, key=loss2.get)

                verse[0]['verstyp'] = verstyp1
                verse[1]['verstyp'] = verstyp2


    def berechne_strophenformel(self, strophen):

        # Initialisierung
        verstypen = [[] for _ in strophen[0]]

        ### Verstypen suchen, welche die Kosten deutlich reduzieren ##########

        for iter in range(2):
            # Kostenreduktion jedes möglichen zusätzlichen Verstyps für jede Position berechnen
            reduction = defaultdict(lambda: defaultdict(float))
            for strophe in strophen:
                for i, versdict in enumerate(strophe):
                    analyses = versdict['analysen']
                
                    # minimale Kosten aller bereits bekannten Verstypen berechnen
                    curr_cost = math.inf
                    for typ, ana in analyses.items():
                        if self.generaltyp(typ) in verstypen[i] and curr_cost > ana[0][0]:
                            curr_cost = ana[0][0]
                        
                    if curr_cost == math.inf:
                        # inf würde hier zu unendlichen Kostenreduktionen führen
                        curr_cost = 3 * max(ana[0][0] for ana in analyses.values())
    
                    # mögliche Kostenreduktionen berechnen
                    for typ, ana in analyses.items():
                        diff = curr_cost - ana[0][0]
                        if diff > 0.0:
                            reduction[i][self.generaltyp(typ)] += diff
    
            if iter == 0:
                verstypen = [[max(red.keys(), key=red.get)]
                             for i, red in reduction.items()]
                continue
            elif not reduction:
                break
            
            # maximale Kostenreduktion für jede Strophenposition berechnen
            # und in Tripeln (i, t, r) speichern
            max_reduction = [(i, *max(red.items(), key=lambda x: x[1]))
                             for i, red in reduction.items()]
        
            # Tripel (i, t, r) nach Kostenreduktion r sortieren
            max_reduction = sorted(max_reduction, key=lambda x: x[2])
    
            # Diejenige Position in der sortierten Liste 
            # mit dem größten Sprung in der Kostenreduktion ermitteln
            _, _, red_values = zip(*max_reduction)
            jumps = [((r2+1) / (r1+1), k)
                     for k, (r1, r2) in enumerate(zip(red_values, red_values[1:]), start=1)]
            max_jump, jump_pos = max(jumps)
    
            # Verstypen oberhalb der Sprungposition in der nach Kostenreduktion
            # sortierten Liste als mögliche Verstypen hinzufügen
            if max_jump > 2.0 and max_reduction[jump_pos][2] > 10:
                for i, t, r in max_reduction[jump_pos:]:
                    verstypen[i].append(t)

            ### statistische Ausgaben zur Information #########################
    
            print('\nKostenreduktion durch zusätzliche Verstypen je Zeile', file=sys.stderr)
            for i, red in sorted(reduction.items()):
                print(i+1, end='  ', file=sys.stderr)
                for t, r in sorted(red.items(), key= lambda x: -x[1])[:8]:
                    print(f'{t} {r:4.0f}', end='  ', file=sys.stderr)
                print('', file=sys.stderr)

        print('\nStrophenformel:', file=sys.stderr)
        for i, typen in enumerate(verstypen):
            print(i+1, ' '.join(typen), sep='  ', file=sys.stderr)

        return verstypen


    def strophen_analyse(self, strophen):
        ''' analysiert eine Folge von Strophen '''

        # Fehlerhafte Eingaben erkennen
        L = len(strophen[0])
        for strophe in strophen:
            if len(strophe) != L:
                print(f'Fehler: Diese Strophe hat eine andere Länge als vorherige Strophen:',
                      file=sys.stderr)
                for versdict in strophe:
                    print(versdict['zeile'], file=sys.stderr)
                return

        if self.strophenformel:
            lf, ls = len(self.strophenformel), len(strophen[0])
            if lf != ls:
                sys.exit(f'Fehler: Länge der Strophenformel ({lf}) passt nicht zur Länge der Strophen ({ls})')
        else:
            if len(strophen) < 2:
                print('Fehler: zu wenige Strophen für eine Strophenanalyse!', file=sys.stderr)
                return
            self.strophenformel = self.berechne_strophenformel(strophen)

        def reimanalyse(verse, formel):
            ''' bestimmt den besten Verstyp '''

            # Ermittle die Menge aller matchenden Kadenzen
            kadenzen = set()
            for vdict in verse:
                for typ in formel[vdict['nummer']]:
                    heb, kad = typ
                    kadenzen.update(['1m', '2m'] if kad == 'm' else
                                    ['2k', '3k'] if kad == 'k' else
                                    ['w'])

            # Suche die optimale Kadenz für alle Reime
            # und den besten damit kompatiblen Verstyp für jeden Vers
            bester_typ = defaultdict(dict)
            for kadenz in kadenzen:
                for i, vdict in enumerate(verse):
                    kosten1, kosten2, kosten3, kosten4,  = [], [], [], []
                    for typ, anas in vdict['analysen'].items():
                        kosten = anas[0][0]
                        kad, heb = typ
                        kadenzmatch = (kad == kadenz)
                        formelmatch = (self.generaltyp(typ) in formel[vdict['nummer']])
                        hebdiff = min((abs(int(heb) - int(t[0])) + 1.5) ** 2
                                      for t in formel[vdict['nummer']])
                        if kadenzmatch and formelmatch:
                            # Bsp.: kad=kadenz=3k und heb=6 und '6k' in formel
                            kosten1.append((kosten, typ))
                        elif formelmatch:
                            # Bsp.: kad=3k und kadenz=*1m und heb=4 und '4k' in formel
                            kosten2.append((kosten+2, typ))
                        elif kadenzmatch:
                            # Bsp.: kad=3k und kadenz=3k und heb=5 und '6k' in formel
                            kosten3.append((kosten+4+hebdiff, typ))
                        else:
                            kosten4.append((kosten+8+hebdiff, typ))
                    bester_typ[kadenz][i] = (min(kosten1) if kosten1 else
                                             min(kosten2) if kosten2 else
                                             min(kosten3) if kosten3 else
                                             min(kosten4))
                
            gesamtkosten = {kadenz: sum(kosten for kosten, _ in versanalysen.values())
                            for kadenz, versanalysen in bester_typ.items()}
            beste_kadenz = min(gesamtkosten.keys(), key=gesamtkosten.get)
        
            for i, (_, typ) in bester_typ[beste_kadenz].items():
                verse[i]['verstyp'] = typ
            ### Ende der Funktion reimanalyse

        for strophe in strophen:
            versgroup = defaultdict(list)
            for i, versdict in enumerate(strophe):
                versdict['nummer'] = i
                versgroup[versdict['index']].append(versdict)
            
            for index, verse in versgroup.items():
                if index == '':
                    for vers in verse:
                        reimanalyse([vers], self.strophenformel)
                else:
                    reimanalyse(verse, self.strophenformel)


    def freie_analyse(self, strophen):
        ''' analysiert schwierigere vorformatierte Text wie den Leich von Walter '''

        def reimanalyse(verse):
            ''' bestimmt den besten Verstyp '''

            # Berechne alle Verstyp-Kandidaten
            kandidaten = set(typ  for vdict in verse  for typ in vdict['analysen'])

            # Berechne die Kosten der häufigsten Verstypen
            gesamtkosten = defaultdict(float)
            for vdict in verse:
                kosten = {typ: anas[0][0]  for typ, anas in vdict['analysen'].items()}
                max_kosten = max(kosten.values())
                for typ in kandidaten:
                    gesamtkosten[typ] += kosten.get(typ, max_kosten * 1.2)

            # Nimm den Verstyp mit den geringsten Kosten und speichere die besten Analysen
            verstyp = min(gesamtkosten, key=gesamtkosten.get)
            for vdict in verse:
                vdict['verstyp'] = verstyp

        for strophe in strophen:
            versgroup = defaultdict(list)
            for i, versdict in enumerate(strophe):
                versdict['nummer'] = i
                versgroup[versdict['index']].append(versdict)
            for index, verse in versgroup.items():
                if index == '':
                    for vers in verse:
                        reimanalyse([vers])
                else:
                    reimanalyse(verse)


    def add_POS_tags(self, strophen):
        all_tokens = [token
                      for strophe in strophen
                      for versdict in strophe
                      for token in versdict['tokens']]
        all_tags = self.tagger(all_tokens)

        # Speicherung der Tags
        for strophe in strophen:
            for versdict in strophe:
                tokens = versdict['tokens']
            
                # Tags der aktuellen Verszeile speichern
                l = len(tokens)
                tags, all_tags = all_tags[:l], all_tags[l:]

                # Interpunktion entfernen
                if len(tokens) != len(tags):
                    sys.exit('Zahl der Tokens und Tags passt nicht in der Zeile: ' +
                             ' '.join(tokens))
                word_tag_pairs = [(token, tag)
                                  for token, tag in zip(tokens, tags)
                                  if token not in self.punctuation]
                if len(word_tag_pairs) == 0:
                    sys.exit('Problem in der Zeile: ' + ' '.join(tokens))
                words, tags = zip(*word_tag_pairs)
                versdict['wörter'] = words
                versdict['tags'] = tags
    

    def analyse(self, filecontent, versart):
        ''' analysiert jeden Vers der Eingabedatei und ruft dann weitere Analysen auf '''

        # Whitespace am Zeilenende entfernen
        text = re.sub(r'[\t ]+\n', '\n', filecontent)
        
        # Text in Strophen und Zeilen aufspalten
        strophen = [[{'zeile': line.rstrip()}
                     for line in strophe.split('\n') if line and not line[0] == '#']
                    for strophe in text.split('\n\n') if strophe.strip()]
        
        for strophe in strophen:
            for versdict in strophe:
                # Indexabtrennung
                if versart in ['frei', 'strophen']:
                    match = re.match(r'([0-9]*)\s*\t\s*(.*)', versdict['zeile'])
                    if not match:
                        sys.exit(f"Fehler in der Zeile: {versdict['zeile']}\n" +
                                 'Erwartetes Format: Index Tabulator Verszeile')
                    versdict['index'] = match.group(1)
                    versdict['zeile'] = match.group(2)

                versdict['norm_zeile'] = self.normalisiere(versdict['zeile'])
                versdict['tokens'] = self.tokenisiere(versdict['norm_zeile'])
            
        # Wortart-Annotation
        self.add_POS_tags(strophen)

        # Silbifizierung
        for strophe in strophen:
            for versdict in strophe:
                versdict['silben'] = syllabify(versdict['wörter'], versdict['tags'])
                versdict['analyseinput'] = self.betonung.make_input(versdict['silben'],
                                                                    versdict['wörter'],
                                                                    versdict['tags'])
                versdict['analysen'] = self.transducer(versdict['analyseinput'], vers_agreement[versart])
                if versdict['analysen'] is None:
                    sys.exit(f"Fehler: Der Transducer liefert keine Analyse für folgende Eingabe:\n\t{versdict['analyseinput']}\nbeim Satz:\n\t{versdict['zeile']}")
        match versart:
            case 'reimpaare':
                self.reimpaar_analyse(strophen)
            case 'strophen':
                self.strophen_analyse(strophen)
            case 'frei':
                self.freie_analyse(strophen)
            case 'default':
                self.default_analyse(strophen)
        self.print_all_analyses(strophen)
        

if __name__ == '__main__':
    args = parse_arguments()
    
    # Lyriktext einlesen und analysieren
    versanalyse = Versanalyse(args)
    with open(args.eingabedatei, encoding='utf-8') as file:
        versanalyse.analyse(file.read(), args.typ)
