import os
import sys
import csv
import calculate
import scipy as sp
import pandas as pd
import scipy.optimize
from pprint import pprint
import matplotlib as mpl
from matplotlib import rcParams
rcParams['font.family'] = 'monospace'
import matplotlib.pyplot as plt
import matplotlib.ticker as mtick
import matplotlib.patches as patches

import warnings
warnings.simplefilter("ignore")

import logging
logger = logging.getLogger('postgres_copy')
logger.addHandler(logging.NullHandler())
logger.propagate = False

mpl.style.use('seaborn-paper')

%matplotlib inline

"""
Explanation: django-postgre-copy speed tests
By Ben Welsh
This notebook tests the effect of dropping database constraints and indexes prior to loading a large data file.
The official PostgreSQL documentation suggests it lead to significant gains.
We will test this claim by dropping constraints and indexes prior to loading data from the California Civic Data Coalition via the django-postgres-copy wrapper on the database's COPY command.
Connect California Civic Data Coalition Django project
Import Python tools
End of explanation
"""


sys.path.insert(0, '/home/palewire/.virtualenvs/django-calaccess-raw-data/src/')
sys.path.insert(0, '/home/palewire/.virtualenvs/django-calaccess-raw-data/lib/python2.7/')
sys.path.insert(0, '/home/palewire/.virtualenvs/django-calaccess-raw-data/lib/python2.7/site-packages/')
sys.path.insert(0, '/home/palewire/Code/django-calaccess-raw-data/')
sys.path.insert(0, '/home/palewire/Code/django-calaccess-raw-data/example/')
sys.path.insert(0, '/home/palewire/Code/django-postgres-copy/')
os.environ.setdefault("DJANGO_SETTINGS_MODULE", "settings")

"""
Explanation: Add the Django settings module to the environment.
End of explanation
"""


import postgres_copy

postgres_copy.__version__

"""
Explanation: Verify we have the correct version of django-postgres-copy
End of explanation
"""


%%capture
import django
django.setup()

"""
Explanation: Boot the Django project
End of explanation
"""


!curl -O https://calaccess.download/latest/raw.zip

!rm *.csv

!unzip raw.zip

"""
Explanation: Prep for speed tests
Download raw data we will test
End of explanation
"""


import calaccess_raw
from calaccess_raw import models 

def truncate(model):
    from django.db import connection
    cursor = connection.cursor()
    cursor.execute('TRUNCATE TABLE "{}";'.format(model._meta.db_table))

# Set a blank list to store the results
result_list = []

# Loop through the models and run the tests
for model in calaccess_raw.get_model_list():
    print "Testing {}".format(model.__name__)

    # Find the CSV for the model
    csv_name = model().get_csv_name()
    if not os.path.exists(csv_name):
        print("No csv. Skipping.")
        continue
    
    # Count its rows
    raw_rows = !wc -l $csv_name
    rows = int(raw_rows[0].split()[0])
    if not rows > 0:
        print("No data. Skipping.")
        continue

    # Map the CSV to the model
    mapping = dict(
        (f.name, f.db_column) for f in model._meta.fields
        if f.db_column
    )
    
    # Set the test function
    f = model.objects.from_csv
    
    # Test it with dropped indexes
    drop_kwargs = dict(drop_constraints=True, drop_indexes=True)
    drop = %%timeit -r 3 -oq truncate(model); f(csv_name, mapping, **drop_kwargs)

    # Test it without dropped indexes
    dont_kwargs = dict(drop_constraints=False, drop_indexes=False)
    dont_drop = %%timeit -r 3 -oq truncate(model); f(csv_name, mapping, **dont_kwargs)

    # Return the results
    result = {
        'model_name': model.__name__,
        'rows': rows,
        'indexed_fields': len(model.objects.indexed_fields),
        'constrained_fields': len(model.objects.constrained_fields),
        'dont_drop': dont_drop.best,
        'drop': drop.best
    }
    pprint(result)
    result_list.append(result)

outfile = csv.DictWriter(
    open("speed-test.csv", 'w'),
    result_list[0].keys()
)
outfile.writeheader()
outfile.writerows(result_list)

"""
Explanation: Import database models
End of explanation
"""


df = pd.read_csv("speed-test.csv")

"""
Explanation: Analyze the results
End of explanation
"""


print "Tables: {:,}".format(len(df))
print "Rows: {:,}".format(df.rows.sum())

"""
Explanation: What we loaded
End of explanation
"""


print "Without drops: {:,.0f} seconds".format(df.dont_drop.sum())
print "With drops: {:,.0f} seconds".format(df['drop'].sum())

"""
Explanation: How long it took in seconds
End of explanation
"""


print "Without drops: {:.0f} minutes, {:.0f} seconds".format(*divmod(df.dont_drop.sum(), 60))
print "With drops: {:.0f} minutes, {:.0f} seconds".format(*divmod(df['drop'].sum(), 60))

"""
Explanation: How long it took in minutes
End of explanation
"""


print "Absolute change: {:.0f} minutes, {:,.0f} seconds".format(*divmod(df['dont_drop'].sum() - df['drop'].sum(), 60))
print "Percent change: {:,.0f}%".format(((df['drop'].sum() - df.dont_drop.sum())/df['dont_drop'].sum())*100)

"""
Explanation: The change
End of explanation
"""


df['change'] = df['dont_drop'] - df['drop']

df['percent_change'] = (df['drop'] - df['dont_drop']) / df['dont_drop']

fig = plt.figure(figsize=(15, 10))
ax = fig.add_subplot(1, 1, 1)
ax.set_xlabel("Row count", fontsize=14, color="#333333")
ax.set_ylabel("Change in load time (seconds)", fontsize=14, color="#333333")
ax.scatter(df.rows, df.change, color="#ba0a35", s=65)
plt.xlim(-500000, plt.xlim()[1]-1000000)
plt.tick_params(
    axis='both',
    which='both',
    bottom='off',
    top='off',
    left="off",
    right="off"
)
ax.spines['right'].set_visible(False)
ax.spines['top'].set_visible(False)
ax.get_xaxis().get_major_formatter().set_scientific(False)
fmt = '{x:,.0f}'
tick = mtick.StrMethodFormatter(fmt)
ax.xaxis.set_major_formatter(tick)
plt.show()

fig = plt.figure(figsize=(15, 10))
ax = fig.add_subplot(1, 1, 1)
ax.set_xlabel("Row count", fontsize=14, color="#333333")
ax.set_ylabel("Percent change in load time", fontsize=14, color="#333333")
ax.scatter(df.rows, df.percent_change, color="#ba0a35", s=65)
plt.xlim(-500000, plt.xlim()[1]-1000000)
plt.axhline(0, lw=0.2, color="#ba0a35")
plt.annotate("FASTER", (8000000, -0.1), fontsize=14, color="#333333")
plt.annotate("SLOWER", (8000000, 0.1), fontsize=14, color="#333333")
ax.add_patch(
    patches.Rectangle(
        (plt.xlim()[0], 0),
        plt.xlim()[1] + 500000,
        plt.ylim()[0],
        color="#ba0a35",
        alpha=0.1
    )
)
plt.tick_params(
    axis='both',
    which='both',
    bottom='off',
    top='off',
    left="off",
    right="off"
)
ax.spines['right'].set_visible(False)
ax.spines['top'].set_visible(False)
ax.get_xaxis().get_major_formatter().set_scientific(False)
fmt = '{x:,.0f}'
tick = mtick.StrMethodFormatter(fmt)
ax.xaxis.set_major_formatter(tick)
plt.show()

"""
Explanation: Charge the changes by table size
End of explanation
"""


