from summertime import model

sample_model = model.summarizer()
documents = [
    """ PG&E stated it scheduled the blackouts in response to forecasts for high winds amid dry conditions. 
    The aim is to reduce the risk of wildfires. Nearly 800 thousand customers were scheduled to be affected 
    by the shutoffs which were expected to last through at least midday tomorrow."""
]
sample_model.summarize(documents)

# ["California's largest electricity provider has turned off power to hundreds of thousands of customers."]

from summertime.model import SUPPORTED_SUMM_MODELS
print(SUPPORTED_SUMM_MODELS)

from summertime import model

# To use a default model
default_model = model.summarizer()    

# Or a specific model
bart_model = model.BartModel()
pegasus_model = model.PegasusModel()
lexrank_model = model.LexRankModel()
textrank_model = model.TextRankModel()

default_model.show_capability()
pegasus_model.show_capability()
textrank_model.show_capability()

documents = [
    """ PG&E stated it scheduled the blackouts in response to forecasts for high winds amid dry conditions. 
    The aim is to reduce the risk of wildfires. Nearly 800 thousand customers were scheduled to be affected 
    by the shutoffs which were expected to last through at least midday tomorrow."""
]

default_model.summarize(documents)
# or 
pegasus_model.summarize(documents)

def __init__(self,
         trained_domain: str=None,
         max_input_length: int=None,
         max_output_length: int=None,
         ):

def summarize(self,
  corpus: Union[List[str], List[List[str]]],
  queries: List[str]=None) -> List[str]:

def show_capability(cls) -> None:

from summertime import dataset

print(dataset.list_all_dataset())

from summertime import dataset

cnn_dataset = dataset.CnndmDataset()
# or 
xsum_dataset = dataset.XsumDataset()
# ..etc

dataset = dataset.CnndmDataset()

train_data = dataset.train_set  
dev_data = dataset.dev_set  
test_data = dataset.test_set        

dataset = dataset.CnndmDataset()

dataset.show_description()

data_instance.source = source    # either `List[str]` or `str`, depending on the dataset itself, string joining may needed to fit into specific models.
data_instance.summary = summary  # a string summary that serves as ground truth
data_instance.query = query      # Optional, applies when a string query is present

print(data_instance)             # to print the data instance in its entirety

data_instance = next(cnn_dataset.train_set)
print(data_instance)

import itertools

# Get a slice from the train set generator - first 5 instances
train_set = itertools.islice(cnn_dataset.train_set, 5)

corpus = [instance.source for instance in train_set]
print(corpus)

from summertime.dataset import CustomDataset

''' The train_set, test_set and validation_set have the following format: 
        List[Dict], list of dictionaries that contain a data instance.
            The dictionary is in the form:
                {"source": "source_data", "summary": "summary_data", "query":"query_data"}
                    * source_data is either of type List[str] or str
                    * summary_data is of type str
                    * query_data is of type str
        The list of dictionaries looks as follows:
            [dictionary_instance_1, dictionary_instance_2, ...]
'''

# Create sample data
train_set = [
    {
        "source": "source1",
        "summary": "summary1",
        "query": "query1",      # only included, if query is present
    }
]
validation_set = [
    {
        "source": "source2",
        "summary": "summary2",
        "query": "query2",      
    }
]
test_set = [
    {
        "source": "source3",
        "summary": "summary3",
        "query": "query3",     
    }
]

# Depending on the dataset properties, you can specify the type of dataset
#   i.e multi_doc, query_based, dialogue_based. If not specified, they default to false
custom_dataset = CustomDataset(
                    train_set=train_set,
                    validation_set=validation_set,
                    test_set=test_set,
                    query_based=True,
                    multi_doc=True
                    dialogue_based=False)

import itertools
from summertime import dataset, model

cnn_dataset = dataset.CnndmDataset()


# Get a slice of the train set - first 5 instances
train_set = itertools.islice(cnn_dataset.train_set, 5)

corpus = [instance.source for instance in train_set]



# Example 1 - traditional non-neural model
# LexRank model
lexrank = model.LexRankModel(corpus)
print(lexrank.show_capability())

lexrank_summary = lexrank.summarize(corpus)
print(lexrank_summary)


# Example 2 - A spaCy pipeline for TextRank (another non-neueral extractive summarization model)
# TextRank model
textrank = model.TextRankModel()
print(textrank.show_capability())

textrank_summary = textrank.summarize(corpus)
print(textrank_summary)


# Example 3 - A neural model to handle large texts
# LongFormer Model
longformer = model.LongFormerModel()
longformer.show_capability()

longformer_summary = longformer.summarize(corpus)
print(longformer_summary)

mbart_model = st_model.MBartModel()
mt5_model = st_model.MT5Model()

# load Spanish portion of MLSum dataset
mlsum = datasets.MlsumDataset(["es"])

corpus = itertools.islice(mlsum.train_set, 5)
corpus = [instance.source for instance in train_set]

# mt5 model will automatically detect Spanish as the language and indicate that this is supported!
mt5_model.summarize(corpus)

from summertime.evaluation import SUPPORTED_EVALUATION_METRICS

print(SUPPORTED_EVALUATION_METRICS)

import summertime.evaluation as st_eval

bert_eval = st_eval.bertscore()
bleu_eval = st_eval.bleu_eval()
meteor_eval = st_eval.bleu_eval()
rouge_eval = st_eval.rouge()
rougewe_eval = st_eval.rougewe()

def __init__(self, metric_name):

def evaluate(self, model, data):

def get_dict(self, keys):

from summertime.evaluation.base_metric import SummMetric
from summertime.evaluation import Rouge, RougeWe, BertScore

import itertools

# Evaluates model on subset of cnn_dailymail
# Get a slice of the train set - first 5 instances
train_set = itertools.islice(cnn_dataset.train_set, 5)

corpus = [instance for instance in train_set]
print(corpus)

articles = [instance.source for instance in corpus]

summaries = sample_model.summarize(articles)
targets = [instance.summary for instance in corpus]

from summertime.evaluation import  BertScore, Rouge, RougeWe,

# Calculate BertScore
bert_metric = BertScore()
bert_score = bert_metric.evaluate(summaries, targets)
print(bert_score)

# Calculate Rouge
rouge_metric = Rouge()
rouge_score = rouge_metric.evaluate(summaries, targets)
print(rouge_score)

# Calculate RougeWe
rougewe_metric = RougeWe()
rougwe_score = rougewe_metric.evaluate(summaries, targets)
print(rougewe_score)

from summertime.pipeline import assemble_model_pipeline
from summertime.dataset import CnndmDataset, QMsumDataset

single_doc_models = assemble_model_pipeline(CnndmDataset)
# [
#   (<model.single_doc.bart_model.BartModel object at 0x7fcd43aa12e0>, 'BART'),
#   (<model.single_doc.lexrank_model.LexRankModel object at 0x7fcd43aa1460>, 'LexRank'),
#   (<model.single_doc.longformer_model.LongformerModel object at 0x7fcd43b17670>, 'Longformer'),
#   (<model.single_doc.pegasus_model.PegasusModel object at 0x7fccb84f2910>, 'Pegasus'),
#   (<model.single_doc.textrank_model.TextRankModel object at 0x7fccb84f2880>, 'TextRank')
# ]

query_based_multi_doc_models = assemble_model_pipeline(QMsumDataset)
# [
#   (<model.query_based.tf_idf_model.TFIDFSummModel object at 0x7fc9c9c81e20>, 'TF-IDF (HMNET)'),
#   (<model.query_based.bm25_model.BM25SummModel object at 0x7fc8b4fa8c10>, 'BM25 (HMNET)')
# ]

# Get test data
import itertools
from summertime.dataset import XsumDataset

# Get a slice of the train set - first 5 instances
sample_dataset = XsumDataset()
sample_data = itertools.islice(sample_dataset.train_set, 100)
generator1 = iter(sample_data)
generator2 = iter(sample_data)

bart_model = BartModel()
pegasus_model = PegasusModel()
models = [bart_model, pegasus_model]
metrics = [metric() for metric in SUPPORTED_EVALUATION_METRICS]

from summertime.evaluation.model_selector import ModelSelector

selector = ModelSelector(models, generator1, metrics)
table = selector.run()
print(table)
visualization = selector.visualize(table)

from summertime.evaluation.model_selector import ModelSelector

new_selector = ModelSelector(models, generator2, metrics)
smart_table = new_selector.run_halving(min_instances=2, factor=2)
print(smart_table)
visualization_smart = selector.visualize(smart_table)

from summertime.evaluation.model_selector import ModelSelector
from summertime.evaluation.error_viz import scatter

keys = ("bert_score_f1", "bleu", "rouge_1_f_score", "rouge_2_f_score", "rouge_l_f_score", "rouge_we_3_f", "meteor")

scatter(models, sample_data, metrics[1:3], keys=keys[1:3], max_instances=5)