#!/usr/bin/env python


import click
import csv
import os
import ujson

from boltons.iterutils import pairwise

from quotes.services import session
from quotes.models import ChadhNovel, BPOArticle, Alignment
from quotes.text import Text
from quotes.utils import open_makedirs


@click.group()
def cli():
    pass


@cli.command()
@click.argument('out_file', type=click.File('w'))
def alignment_csv(out_file):
    """Dump alignments to CSV.
    """
    writer = csv.DictWriter(out_file, (

        'chadh_slug',
        'bpo_pub_title',
        'bpo_title',
        'bpo_article_type',

        'a_start',
        'b_start',
        'size',
        'a_token_count',

        'a_prefix',
        'a_snippet',
        'a_suffix',

        'b_prefix',
        'b_snippet',
        'b_suffix',
    ))

    writer.writeheader()

    alignments = (
        session
        .query(Alignment)
        .join(ChadhNovel, BPOArticle)
        .filter(Alignment.size >= 5)
        .yield_per(1000)
    )

    counts = {}

    for i, a in enumerate(alignments):

        slug = a.chadh_novel.slug

        if slug not in counts:
            text = Text(a.chadh_novel.text)
            counts[slug] = len(text.tokens)

        writer.writerow(dict(

            chadh_slug=a.chadh_novel.slug,
            bpo_pub_title=a.bpo_article.publication_title,
            bpo_title=a.bpo_article.record_title,
            bpo_article_type=a.bpo_article.object_type,

            a_start=a.a_start,
            b_start=a.b_start,
            size=a.size,
            a_token_count=counts[slug],

            a_prefix=a.a_prefix,
            a_snippet=a.a_snippet,
            a_suffix=a.a_suffix,

            b_prefix=a.b_prefix,
            b_snippet=a.b_snippet,
            b_suffix=a.b_suffix,

        ))

        if i % 1000 == 0:
            print(i)


@cli.command()
@click.argument('out_dir', type=click.Path(exists=True))
@click.option('--min_wc', default=50e6)
def dump_slices(out_dir, min_wc):
    """Dump time-sliced text JSON, for experiments with historical vectors.
    """
    for y1, y2 in pairwise(range(1700, 2000, 20)):

        # Select articles in slice.
        articles = (
            BPOArticle.query
            .filter(BPOArticle.year >= y1, BPOArticle.year <= y2)
            .order_by(BPOArticle.record_id)
        )

        # Count words in articles, until min WC.
        wc = 0
        for i, article in enumerate(articles):
            if article.text:
                wc += len(article.text.split(' '))
                if wc > min_wc:
                    break

        # Skip slices without enough words.
        if wc < min_wc:
            continue

        # Dump JSON.
        for article in articles.limit(i):

            fname = str(article.record_id)+'.json'
            path = os.path.join(out_dir, str(y1), fname)

            with open_makedirs(path, 'w') as fh:
                ujson.dump(dict(article), fh)

        print(y1, y2, i, wc)


@cli.command()
@click.argument('out_dir', type=click.Path(exists=True))
def dump_blackwood(out_dir):
    """Dump Blackwood articles.
    """
    # Select Blackwood articles.
    articles = (
        BPOArticle.query
        .filter_by(publication_title="Blackwood's Edinburgh magazine")
    )

    # Dump JSON.
    for article in articles:

        fname = str(article.record_id)+'.json'
        path = os.path.join(out_dir, fname)

        with open_makedirs(path, 'w') as fh:
            ujson.dump(dict(article), fh)


if __name__ == '__main__':
    cli()
