import os
os.system("python -m spacy download it_core_news_sm")

import gradio as gr

from transformers import AutoTokenizer, PreTrainedTokenizerFast
from transformers import AutoModelForTokenClassification
from transformers import pipeline

import spacy
from spacy import displacy
from spacy.tokens import Span

background_colors_entity_tag = {
    "RIC": "#ff5e5e",
    "RCR": "#ff9999",
    "CTR": "#ffd699",
    "AVV": "#80c5c5",
    "CNS": "#ff9500",
    "PMI": "#0ea5e9",
    "DOM": "#c3a1c9",
    "CDA": "#84b351",
    "SNT": "#ffff5e",
}

css = {
'entity_tag': 'color:#000;background: #xxxxxx; font-size: 0.8em; font-weight: bold; line-height: 2.5; border-radius: 0.35em; text-transform: uppercase; vertical-align: middle; margin-left: 0.5em;'
}
entity_list = "<span style='"
entity_list += f"{css['entity_tag'].replace('#xxxxxx',background_colors_entity_tag['RIC'])};padding:0.5em;"
entity_list += "'>RICORSO</span>"
entity_list += "<span style='"
entity_list += f"{css['entity_tag'].replace('#xxxxxx',background_colors_entity_tag['RCR'])};padding:0.5em;"
entity_list += "'>RICORRENTE</span>"
entity_list += "<span style='"
entity_list += f"{css['entity_tag'].replace('#xxxxxx',background_colors_entity_tag['CTR'])};padding:0.5em;"
entity_list += "'>CONTRORICORRENTE</span>"
entity_list += "<span style='"
entity_list += f"{css['entity_tag'].replace('#xxxxxx',background_colors_entity_tag['AVV'])};padding:0.5em;"
entity_list += "'>AVVOCATO</span>"
entity_list += "<span style='"
entity_list += f"{css['entity_tag'].replace('#xxxxxx',background_colors_entity_tag['CNS'])};padding:0.5em;"
entity_list += "'>CONSIGLIERE</span>"
entity_list += "<span style='"
entity_list += f"{css['entity_tag'].replace('#xxxxxx',background_colors_entity_tag['PMI'])};padding:0.5em;"
entity_list += "'>PUBBLICO MINISTERO</span>"
entity_list += "<span style='"
entity_list += f"{css['entity_tag'].replace('#xxxxxx',background_colors_entity_tag['DOM'])};padding:0.5em;"
entity_list += "'>DOMICILIO</span>"
entity_list += "<span style='"
entity_list += f"{css['entity_tag'].replace('#xxxxxx',background_colors_entity_tag['CDA'])};padding:0.5em;"
entity_list += "'>CORTE D'APPELLO</span>"
entity_list += "<span style='"
entity_list += f"{css['entity_tag'].replace('#xxxxxx',background_colors_entity_tag['SNT'])};padding:0.5em;"
entity_list += "'>SENTENZA</span>"

title = "NER for Italian Legal Documents"
description = "This application shows a Named Entity Recognition (NER) model capable of extracting entities of interest (e.g. claimant, defendant, lawyer and court) from Italian civil judgments issued by the Corte Suprema di Cassazione."
article = f"<div>Named Entities: {entity_list}</div>"

examples = [
            """la seguente SENTENZA sul ricorso 24817-2015 proposto da: ANDREA FORMISANO, elettivamente domiciliato in ROMA VIA S. TOMMASO D'AQUINO 7, presso lo studio dell'avvocato CARLO BORELLO, che lo rappresenta e difende giusta delega in calce; - ricorrente - contro SOGET SPA, CAMERA DI COMMERCIO DI PESCARA; - intimati - avverso la sentenza n. 169/2012 della COMM.TRIB.REG.SEZ.DIST. di PESCARA, depositata il 13/03/2012; udita la relazione della causa svolta nella pubblica udienza del 04/04/2018 dal Consigliere Dott. MILENA BALSAMO; udito il P.M. in persona del Sostituto Procuratore Generale Dott. GIOVANNI GIACALONE che ha concluso per l'inammissibilità in subordine rigetto del ricorso.""",
            """la seguente SENTENZA sul ricorso 17668-2016 proposto da: C.B.H. CITTA DI BARI HOSPITAL S.P.A., in persona del legale rappresentante pro tempore, elettivamente domiciliata in ROMA, LUNGOTEVERE DEI MELLINI 10, presso lo studio dell'avvocato CRISTIANO MARINESE, rappresentata e difesa dagli avvocati GIUSEPPE LUIGI 2022 POLITO, FRANCESCO ANTONUCCI; 51 - ricorrente - contro I.N.P.S. - ISTITUTO NAZIONALE PREVIDENZA SOCIALE, in persona del legale rappresentante pro tempore, elettivamente domiciliato in ROMA, VIA CESARE BECCARIA 29, presso l'Avvocatura Centrale dell'Istituto, rappresentato e difeso dagli avvocati ANTONINO SGROI, CARLA D'ALOISIO, ESTER ADA SCIPLINO, EMANUELE DE ROSE, LELIO MARITATO, GIUSEPPE MATANO; - controricorrente - nonchè contro EQUITALIA SERVIZI DI RISCOSIONE S.P.A. già EQUITALIA SUD S.P.A. agente della riscossione della provincia di Bari; - intimata - avverso la sentenza n. 2696/2015 della CORTE D'APPELLO di BARI, depositata il 13/01/2016 R.G.N. 1439/2013; udita la relazione della causa svolta nella pubblica udienza del 12/01/2022 dal Consigliere Dott. DANIELA CALAFIORE; udito il P.M. in persona del Sostituto Procuratore Generale Dott. STEFANO VISONA' che ha concluso per il rigetto del ricorso; udito l'avvocato ANTONINO SGROI. R.g. n. 17668/2016""",
            """4. SENTENZA sul ricorso 4005-2012 proposto da: BANCA NAZIONALE DEL LAVORO S.P.A. C.E. 09339391006, in persona del legale rappresentante pro tempore, elettivamente domiciliata in ROMA, VIA PO 25/B, presso lo studio degli avvocati ROBERTO PESSI, FRANCESCO GIAMMARIA, che la rappresentano e difendono, giusta procura speciale notarile in atti; 2015 - ricorrente - 4680 contro CAMPAGNOLI ALESSANDRO MARIA C.F. CMPLSN59L29G388P; 4 - intimato - Nonché da: CAMPAGNOLI ALESSANDRO MARIA C.E. CMPLSN59L29G388P, domiciliato in ROMA PIAZZA CAVOUR, presso LA CANCELLERIA DELLA CORTE SUPREMA DI CASSAZIONE, rappresentato e difeso dall'avvocato FABRIZIA MAURICI, giusta procura speciale notarile in atti; - controricorrente e ricorrente incidentale - contro BANCA NAZIONALE DEL LAVORO S.P.A. C.E. 09339391006, in persona del legale rappresentante pro tempore, elettivamente domiciliata in ROMA, VIA PO 25/B, presso lo studio degli avvocati ROBERTO FESSI, FRANCESCO GIAMMARIA, che la rappresentano e difendono, giusta procura speciale notarile in atti; - controricorrente al ricorso incidentale - avverso la sentenza n. 1091/2011 della CORTE D'APPELLO di MILANO, depositata il 28/10/2011 R.G.N. 537/2008; udita la relazione della causa svolta nella pubblica udienza del 02/12/2015 dal Consigliere Dott. UMBERTO BERRINO; udito l'Avvocato SERRANI TIZIANA per delega verbale FESSI ROBERTO; udito l'Avvocato MAURICI FABRIZIA (per procura speciale notarile); udito il P.M. in persona del Sostituto Procuratore Generale Dott. RITA SANLORENZO che ha concluso per il rigetto del ricorso principale e del ricorso incidentale.""",
            #"""SENTENZA sul ricorso 11948-2014 proposto da: VENTURA VINCENZO C.F. VNTVCN47T08A841S, già elettivamente domiciliato in ROMA, VIA VALLISNERI 11, presso lo studio dell'avvocato PAOLO PACIFICI, che lo rappresenta e difende unitamente all'avvocato DIEGO TOSI, giusta delega in atti e da ultimo domiciliato 2015 presso LA CANCELLERIA DELLA CORTE SUPREMA DI CASSAZIONE; 4525 - ricorrente - contro k RAI RADIOTELEVISIONE ITALIANA S.P.A. C.F. 06382641006, in persona del legale rappresentante pro tempore, elettivamente domiciliata in ROMA, VIA P.L. DA PALESTRINA 47, presso lo studio dell'avvocato RINALDO GEREMIA, rappresentata e difesa dall'avvocato NATALIA FERRO, giusta delega in atti; - controri corrente nonchè contro I.N.A.I.L - ISTITUTO NAZIONALE PER L'ASSICURAZIONE CONTRO GLI INFORTUNI SUL LAVORO C.F. 01165400589, in persona del legale rappresentante pro tempore, elettivamente domiciliato in ROMA, VIA IV NOVEMBRE 144, presso lo studio degli avvocati LUCIANA ROMEO, LETIZIA CRIPPA, che lo rappresentano e difendono giusta delega in atti; - controricorrente - avverso la sentenza n. 1423/2013 della CORTE D'APPELLO di TORINO, depositata il 03/02/2014 R.G.N. 275/2013; udita la relazione della causa svolta nella pubblica udienza del 25/11/2015 dal Consigliere Dott. NICOLA DE MARINIS; AVV, udito l'Avvocato OTTOLINI TERESA per delega', ROMEO LUCIANA; udito l'Avvocato GEREMIA RINALDO per delega'-eFERRO NATALIA; udito il P.M. in persona del Sostituto Procuratore Generale Dott. RENATO FINOCCHI GHERSI che ha concluso per ESTINZIONE PER RINUNCIA. ... , z , I ? F""",

]

model_name = "fabiod20/italian-legal-ner"
model = AutoModelForTokenClassification.from_pretrained(model_name, use_auth_token=os.environ['token'])
tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=os.environ['token'])

ner_pipe = pipeline("ner", model=model, tokenizer=tokenizer)

nlp = spacy.load("it_core_news_sm")
nlp.disable_pipes("ner")

def ner(input_text):
  entities = ner_pipe(input_text, aggregation_strategy="first")

  doc = nlp(input_text)

  potential_entities = []

  for entity in entities:
    start = entity["start"]
    end = entity["end"]
    label = entity["entity_group"]

    ent = doc.char_span(start, end, label=label)
    if ent != None:
      doc.ents += (ent,)
    else:
      potential_entities.append(entity)

  potential_entities.append({"entity_group": "NONE", "start": -1, "end": -1})

  start = potential_entities[0]["start"]
  end = potential_entities[0]["end"]
  label = potential_entities[0]["entity_group"]

  for item in potential_entities:
    if item["entity_group"] == label and item["start"] == end:
      end = item["end"]
      continue
    else:
      if item["start"] != start:
        ent = doc.char_span(start, end, label=label)
        doc.ents += (ent,)

      start = item["start"]
      end = item["end"]
      label = item["entity_group"]

  colors = {
    "RIC": "#ff5e5e",
    "RCR": "#ff9999",
    "CTR": "#ffd699",
    "DOM": "#c3a1c9",
    "AVV": "#80c5c5",
    "CNS": "#ff9500",
    "PMI": "#0ea5e9",
    "CDA": "#84b351",
    "SNT": "#ffff5e",
  }
  options = {"ents": colors.keys(), "colors": colors}

  output = displacy.render(doc, style="ent", options=options)
  return output

interface = gr.Interface(
    title=title,
    description=description,
    article=article,
    allow_screenshot=False,
    allow_flagging=False,
    fn=ner,
    inputs=gr.inputs.Textbox(placeholder="Insert an Italian judgments (you can click on an example below)", lines=10),
    outputs=gr.outputs.HTML(),
    examples=examples
    )

interface.launch()