import os

import luigi
import scipy as sp

from connectome.model import EXP
from figures import OutputFolder, style
from figures.analysispipeline.auxiliary import plot
import figures.style.linestyle
import figures.style.size

connectivity_over_distance_scale = .5
y_name = "Conn.\nprob. "

def break_maybe(s):
    if s[0] == "$":
        return s
    return s + "\n"

y_label_short = style.latex_name("exc_and_inh_connectivity")
y_label_long = y_name

from ..auxiliary import FolderPlot


def common_style(ax, y_label):
    ax.set_xlabel(style.latex_name("relative_intersoma_distance"))
    ax.set_ylabel(y_label)
    style.reduced_ticks_0_1(ax, label_pad_reduce=-2)
    ax.yaxis.labelpad = 0



class ExpConnectivityOverDistance(FolderPlot):
    SCALE = style.size.xsq

    def run(self):
        barrel_size = sp.linspace(0, 1)
        exp = EXP(nr_neurons=2000, inh_ratio=.1, p_exc=.2, p_inh=.6, decay=1)
        exp()
        exc_label = style.latex_name("excitatory_decay_length") + " ={:.1f}$".format(exp.alpha_exc_)
        exc_p = sp.exp(-barrel_size/exp.alpha_exc_)
        inh_label = style.latex_name("inhibitory_decay_length") + " ={:.1f}$".format(exp.alpha_inh_)
        inh_p = sp.exp(-barrel_size/exp.alpha_inh_)
        middle_index = len(exc_p)//2
        offset = .05
        with self.plot("exp_connectivity_over_distance", scale=self.SCALE) as (fig, ax):
            ax.plot(barrel_size, exc_p, "-", color="k")
            ax.plot(barrel_size, inh_p, dashes=style.linestyle.dashed, color="k")
            ax.text(barrel_size[middle_index], exc_p[middle_index] + offset, exc_label)
            ax.text(barrel_size[middle_index], inh_p[middle_index] + offset -.02, inh_label)
            common_style(ax, y_label_short)


class ERConnectivityOverDistance(FolderPlot):
    def run(self):
        with self.plot("er_connectivity_over_distance", scale=ExpConnectivityOverDistance.SCALE) as (fig, ax):
            ax.plot([0, 1], [.2] * 2, label=style.latex_name("p_exc"), color="k")
            ax.plot([0, 1], [.6] * 2, dashes=style.linestyle.dashed, label=style.latex_name("p_inh"), color="k")
            ax.text(1.05, .2, style.latex_name("p_ee"), va="center", ha="left")
            ax.text(1.05, .6, style.latex_name("p_ii"), va="center", ha="left")
            common_style(ax, y_label_long)