# -*- coding: utf-8 -*-
# Generated by Django 1.11 on 2017-09-02 20:29
from __future__ import unicode_literals
import os

from graphviz import Digraph

from django.db import migrations
from django.core.files.base import ContentFile
from django.contrib.contenttypes.management import create_contenttypes

from bayesian_networks.bayespy_constants import (
    DIST_GAUSSIAN_ARD, DIST_GAMMA)
from bayesian_networks.models import BayesianNetworkNode as BNN


def generate_bn_image(bn):
    """
    Auxiliary function for generating the image of a BN, as model
    methods are not available in migrations.
    """
    dot = Digraph(comment=bn.name)
    nodes = bn.nodes.all()
    for node in nodes:
        dot.node(name=node.name, label=node.name)
    edges = bn.edges.all()
    for edge in edges:
         dot.edge(str(edge.parent.name),
                 str(edge.child.name))
    dot.format = "png"
    contentfile = ContentFile(dot.pipe())
    image_name = "{0}/{1}".format(
            os.path.join("django_ai",
                         "bayesian_networks"),
            bn.name + ".png")
    bn.image.save(image_name, contentfile)
    bn.save()


def create_bn1_example(apps, schema_editor):
    # Content Types Hackery for ensuring that it exists
    app_config = apps.get_app_config('examples')
    app_config.models_module = app_config.models_module or True
    create_contenttypes(app_config)
    ##

    BayesianNetwork = apps.get_model("bayesian_networks",
                                     "BayesianNetwork")
    BayesianNetworkEdge = apps.get_model("bayesian_networks",
                                         "BayesianNetworkEdge")
    BayesianNetworkNode = apps.get_model("bayesian_networks",
                                         "BayesianNetworkNode")
    ContentType = apps.get_model("contenttypes",
                                "ContentType")
    
    bn1 = BayesianNetwork(name="BN1 (Example)")
    bn1.save()
    mu = BayesianNetworkNode(
        network=bn1,
        name="mu",
        node_type=BNN.NODE_TYPE_STOCHASTIC,
        is_observable=False,
        distribution=DIST_GAUSSIAN_ARD,
        distribution_params="0, 1e-6",
        graph_interval="-10, 20"
    )
    tau = BayesianNetworkNode(
        network=bn1,
        name="tau",
        node_type=BNN.NODE_TYPE_STOCHASTIC,
        is_observable=False,
        distribution=DIST_GAMMA,
        distribution_params="1e-6, 1e-6",
        graph_interval="1e-6, 0.1"
    )
    ui_avg1 = BayesianNetworkNode(
        network=bn1,
        name="userinfo.avg1",
        node_type=BNN.NODE_TYPE_STOCHASTIC,
        is_observable=True,
        distribution=DIST_GAUSSIAN_ARD,
        distribution_params="mu, tau",
        ref_model=ContentType.objects.get(model="userinfo",
                                          app_label="examples"),
        ref_column="avg1",

    )
    mu.save()
    tau.save()
    ui_avg1.save()

    mu_to_ui_avg1 = BayesianNetworkEdge(
        network=bn1,
        description="mu -> userinfo.avg1",
        parent=mu,
        child=ui_avg1
    )
    tau_to_ui_avg1 = BayesianNetworkEdge(
        network=bn1,
        description="tau -> userinfo.avg1",
        parent=tau,
        child=ui_avg1
    )
    mu_to_ui_avg1.save()
    tau_to_ui_avg1.save()
    # Generate the image
    generate_bn_image(bn1)


def delete_bn1_example(apps, schema_editor):
    BayesianNetwork = apps.get_model("bayesian_networks",
                                     "BayesianNetwork")

    BayesianNetwork.objects.get(name="BN1 (Example)").delete()


class Migration(migrations.Migration):

    dependencies = [
        ('examples', '0003_populate_userinfo'),
    ]

    operations = [
        migrations.RunPython(create_bn1_example,
                             delete_bn1_example),
    ]