import pandas as pd

# load data
exprs = pd.read_table(snakemake.input[0], index_col=[0, 1])
exprs = (exprs["corrected"] + exprs["exact"]).unstack(0).fillna(0)
exprs.to_csv(snakemake.output[0], sep="\t")

import pandas as pd

# load data
exprs = pd.read_table(snakemake.input[0], index_col=[0, 1])["expr_map"].unstack(0).fillna(0)
exprs.to_csv(snakemake.output[0], sep="\t")

import svgutils.transform as sg
from common import load_svg, label_plot

fx, fy = "5.1in", "1.3in"
if snakemake.wildcards.dataset == "1001genesData":
    fx, fy = "5.5in", "1.3in"
fig = sg.SVGFigure(fx, fy)
a = load_svg(snakemake.input.a)
b = load_svg(snakemake.input.b)
c = load_svg(snakemake.input.c)
d = load_svg(snakemake.input.d)

xb = 170
xc = 275
xd = 380
if snakemake.wildcards.dataset == "1001genesData":
    xb = 170
    xc = 290
    xd = 405

b.moveto(xb, 0)
c.moveto(xc, 0)
d.moveto(xd, 0)

la = label_plot(5,10, "a")
lb = label_plot(xb,10, "b")
lc = label_plot(xc,10, "c")
ld = label_plot(xd,10, "d")

fig.append([a, b, c, d, la, lb, lc, ld])
fig.save(snakemake.output[0])

import svgutils.transform as sg
from common import load_svg, label_plot

fig = sg.SVGFigure("7.5in", "1.9in")
a = load_svg(snakemake.input.mhd4[0])
b = load_svg(snakemake.input.mhd4[1])
c = load_svg(snakemake.input.mhd2[0])
d = load_svg(snakemake.input.mhd2[1])
b.moveto(170, 0)
c.moveto(340, 0)
d.moveto(510, 0)

la = label_plot(5, 10, "a")
lb = label_plot(175, 10, "b")
lc = label_plot(345, 10, "c")
ld = label_plot(515, 10, "d")

fig.append([a, b, c, d, la, lb, lc, ld])
fig.save(snakemake.output[0])

import svgutils.transform as sg
from common import load_svg, label_plot


fig = sg.SVGFigure("4.1in", "1.8in")
a = load_svg(snakemake.input[1])
b = load_svg(snakemake.input[0])
b.moveto(190, 0)

la = label_plot(5, 10, "a")
lb = label_plot(185, 10, "b")

fig.append([a, b, la, lb])
fig.save(snakemake.output[0])

import svgutils.transform as sg
from common import load_svg, label_plot


fig = sg.SVGFigure("5.3in", "1.8in")
a = load_svg(snakemake.input.a)
b = load_svg(snakemake.input.b)
a.moveto(10, 0)
b.moveto(260, 0)

la = label_plot(5, 10, "a")
lb = label_plot(255, 10, "b")

fig.append([a, b, la, lb])
fig.save(snakemake.output[0])

import svgutils.transform as sg
from common import load_svg, label_plot


fig = sg.SVGFigure("5.8in", "1.8in")
a = load_svg(snakemake.input.a)
b = load_svg(snakemake.input.b)
b.moveto(260, 0)

la = label_plot(5, 10, "a")
lb = label_plot(255, 10, "b")

fig.append([a, b, la, lb])
fig.save(snakemake.output[0])

import svgutils.transform as sg
from common import load_svg, label_plot

fig = sg.SVGFigure("6.4in", "1.7in")
a = load_svg(snakemake.input[0])
b = load_svg(snakemake.input[1])
#c = load_svg(snakemake.input.c)
#d = load_svg(snakemake.input.d)
e = load_svg(snakemake.input[2])
#f = load_svg(snakemake.input.f)
b.moveto(190, 0)
#c.moveto(380, 0)
#d.moveto(0, 160)
e.moveto(380, 0)
#f.moveto(380, 160)

la = label_plot(5,10, "a")
lb = label_plot(195,10, "b")
lc = label_plot(385,10, "c")

fig.append([a, b, e, la, lb, lc])
fig.save(snakemake.output[0])

import svgutils.transform as sg
from common import load_svg, label_plot

fig = sg.SVGFigure("6.5in", "2.1in")
a = load_svg(snakemake.input[0])
b = load_svg(snakemake.input[1])
c = load_svg(snakemake.input[2])
d = load_svg(snakemake.input[3])
a.moveto(10, 10, scale=0.6)
b.moveto(20, 130, scale=0.4)
c.moveto(250, 10, scale=0.65)
d.moveto(440, 10, scale=1.0)


la = label_plot(5, 10, "a")
lb = label_plot(5, 130, "b")
lc = label_plot(240, 10, "c")
ld = label_plot(430, 10, "d")

fig.append([a, b, c, d, la, lb, lc, ld])
fig.save(snakemake.output[0])

import svgutils.transform as sg
from common import load_svg, label_plot


fig = sg.SVGFigure("6.5in", "1.8in")
a = load_svg(snakemake.input.a)
b = load_svg(snakemake.input.b)
b.moveto(230, 0)

la = label_plot(5, 10, "a")
lb = label_plot(225, 10, "b")

fig.append([a, b, la, lb])
fig.save(snakemake.output[0])

import svgutils.transform as sg
from common import load_svg, label_plot


fig = sg.SVGFigure("5.8in", "3.3in")
a = load_svg(snakemake.input.a)
b = load_svg(snakemake.input.b)
c = load_svg(snakemake.input.c)
d = load_svg(snakemake.input.d)
e = load_svg(snakemake.input.e)
f = load_svg(snakemake.input.f)
b.moveto(260, 0)
c.moveto(0, 160)
d.moveto(128, 160)
e.moveto(250, 160)
f.moveto(384, 160)

la = label_plot(5, 10, "a")
lb = label_plot(265, 10, "d")
lc = label_plot(5, 170, "b")
ld = label_plot(130, 170, "c")
le = label_plot(255, 170, "e")
lf = label_plot(385, 170, "f")

fig.append([a, b, c, d, e, f, la, lb, lc, ld, le, lf])
fig.save(snakemake.output[0])

import svgutils.transform as sg
from common import load_svg, label_plot


fig = sg.SVGFigure("5.6in", "1.8in")
a = load_svg(snakemake.input.a)
b = load_svg(snakemake.input.b)
a.moveto(10, 0)
b.moveto(260, 0)

la = label_plot(5, 10, "a")
lb = label_plot(255, 10, "b")

fig.append([a, b, la, lb])
fig.save(snakemake.output[0])

import pandas as pd

cdf = pd.read_table(snakemake.input.cdf, index_col=0)
scales = pd.read_table(snakemake.input.scales, index_col=0, squeeze=True, header=None)
cdf["expr"] *= scales[int(snakemake.wildcards.experiment)]
cdf.to_csv(snakemake.output[0], sep="\t")

import pandas as pd

expr = pd.read_table(snakemake.input.expr, index_col=0)
scales = pd.read_table(snakemake.input.scales, index_col=0, squeeze=True, header=None)
expr *= scales[int(snakemake.wildcards.experiment)]
expr.to_csv(snakemake.output[0], sep="\t")

import pandas as pd
import numpy as np


exprs = [pd.read_table(f, index_col=0) for f in snakemake.input]

# calculate upper quartiles
quartiles = np.array([e.unstack().quantile(0.75) for e in exprs])
mean_quartile = quartiles.mean()
# calculate scale factors for upper quartile normalization
scales = pd.Series([mean_quartile / q for q in quartiles], index=snakemake.params.experiments, name="scale_factors")
scales.to_csv(snakemake.output[0], sep="\t")

import matplotlib
matplotlib.use("agg")
import matplotlib.pyplot as plt
import seaborn as sns
import pandas as pd
import numpy as np


PSEUDOCOUNTS = 1


sns.set(style="ticks", palette="colorblind", context=snakemake.wildcards.context)
plt.figure(figsize=[snakemake.config["plots"]["figsize"][1]] * 2)


def get_means(raw_counts):
    return raw_counts.groupby(level=1).mean() + PSEUDOCOUNTS


def get_cvs(batch_means):
    g = batch_means.groupby(level=1)
    return g.std() / g.mean()


def read_raw_counts(f):
    m = pd.read_table(f, index_col=[0, 1]).sum(axis=1).unstack(fill_value=0)
    return m.stack()


def normalize(raw_counts):
    quartiles = np.array([counts.quantile(0.75) for counts in raw_counts])
    mean_quartile = quartiles.mean()
    print([mean_quartile / q for counts, q in zip(raw_counts, quartiles)])
    return [counts * (mean_quartile / q) for counts, q in zip(raw_counts, quartiles)]

raw_counts = normalize([read_raw_counts(f) for f in snakemake.input.raw_counts])
batch_means = [get_means(counts) for counts in raw_counts]
batch_means = pd.concat(batch_means, keys=range(len(batch_means)))

raw_cvs = get_cvs(batch_means)

#print(raw_cvs[raw_cvs < 0.1])

posterior_cvs = pd.read_table(snakemake.input.diffexp, index_col=0)[snakemake.wildcards.estimate]

plt.scatter(posterior_cvs, raw_cvs.loc[posterior_cvs.index], s=2, c="k", alpha=0.6, edgecolors="face")
#sns.regplot(posterior_cvs, raw_cvs.loc[posterior_cvs.index], color="k", scatter_kws={"s": 2, "alpha": 0.6})
limits = (0, min(plt.xlim()[1], plt.ylim()[1]))
limits = (0, 1)

plt.plot(limits, limits, "r--")

plt.xlim(limits)
plt.ylim(limits)
plt.xlabel("posterior CV" + (" (conservative)" if snakemake.wildcards.estimate == "cv_ci_lower" else ""))
plt.ylabel("raw CV")
sns.despine()
plt.savefig(snakemake.output[0], bbox_inches="tight")

print(raw_cvs)

import pandas as pd
import matplotlib as mpl
mpl.use("agg")
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn import metrics
import numpy as np

sns.set(style="ticks", palette="colorblind", context=snakemake.wildcards.context)

small = pd.concat([pd.read_table(f, index_col=0) for f in snakemake.input.small], axis=1)
large = pd.concat([pd.read_table(f, index_col=0) for f in snakemake.input.large], axis=1)

common_genes = small.index.intersection(large.index)
means = lambda estimates: estimates.loc[common_genes].mean(axis="columns")

small = means(small)
large = means(large)
smape = (large - small).abs().sum() / (large + small).sum()
#relative_accuracy = np.log10(large / small)
small = np.log10(small)
large = np.log10(large)

min_value = np.log10(0.1)#min(small.min(), large.min())
max_value = np.log10(100)#max(small.max(), large.max())

fig = plt.figure(figsize=snakemake.config["plots"]["figsize"])
fig.add_subplot(111, aspect='equal')

sns.regplot(small, large, line_kws={"color": "red"}, scatter_kws={"color": "black", "s": 3})
plt.plot([min_value, max_value], [min_value, max_value], "k--")

label = "raw counts" if snakemake.wildcards.type == "counts" else "posterior counts"
plt.xlabel("MHD4 " + label)
plt.ylabel("MHD2 " + label)
plt.xlim((min_value, max_value))
plt.ylim((min_value, max_value))
plt.text(x=0.5, y=1, s="SMAPE={:.0%}".format(smape), horizontalalignment="center", verticalalignment="top", transform=plt.gca().transAxes)

logticks = lambda ticks: ["$10^{{{:.0f}}}$".format(y) for y in ticks]
ax = plt.gca()
ax.locator_params(nbins=4)
ax.set_xticklabels(logticks(ax.get_xticks()))
ax.set_yticklabels(logticks(ax.get_yticks()))
sns.despine()

plt.savefig(snakemake.output[0], bbox_inches="tight")

import numpy as np
import pandas as pd
import matplotlib as mpl
mpl.use("svg")
import matplotlib.pyplot as plt
import seaborn as sns

import re

pattern = re.compile(r"DEBUG: TYPE=EM-iteration, CELL=(?P<cell>\d+), x=(?P<expr>\[(\d+, )+\d+\])")


exprs = []

with open(snakemake.input[0]) as f:
    for l in f:
        m = pattern.match(l)
        if re.match(r"DEBUG: TYPE=EM-iteration, CELL=(?P<cell>\d+)", l):
            print("match")
        if m and m.group("cell") == "1":
            exprs.append(eval(m.group("expr")))

exprs = pd.DataFrame(exprs)


sns.set(style="ticks", palette="colorblind", context=snakemake.wildcards.context)
x, y = snakemake.config["plots"]["figsize"]
plt.figure(figsize=(x * 1.5, y))

for feat, e in exprs.iteritems():
    curr = e[1:].reset_index(drop=True)
    last = e[:-1].reset_index(drop=True)
    #e = ((curr + 1) / (last + 1))
    plt.semilogy(np.arange(e.size), e, "k-", linewidth=1, alpha=0.5)

sns.despine()

plt.xlabel("EM iteration")
plt.ylabel("expression")
#plt.ylim((0, 200))
plt.savefig(snakemake.output[0], bbox_inches="tight")

import numpy as np
import pandas as pd
import matplotlib as mpl
mpl.use("svg")
import matplotlib.pyplot as plt
import seaborn as sns

sns.set(style="ticks", palette="colorblind", context=snakemake.wildcards.context)
x, y = snakemake.config["plots"]["figsize"]
plt.figure(figsize=(x * 1.5, y))


error_rates = pd.concat([pd.read_table(f, index_col=0) for f in snakemake.input],
                        keys=snakemake.params.experiments)

error_rates = pd.DataFrame(error_rates.stack()).reset_index()
error_rates.columns = ["experiment", "pos", "error-rate", "rate"]
error_rates["pos"] += 1

sns.barplot(x="pos", y="rate", hue="error-rate", data=error_rates, errwidth=1, linewidth=0, palette=["red", "grey"], saturation=0.7)
plt.ylim((0.0, plt.ylim()[1] + 0.1))
plt.ylabel("")
plt.xlabel("position")
sns.despine()
plt.savefig(snakemake.output[0], bbox_inches="tight")

print("mean error rates per position")
means = error_rates.groupby(["error-rate", "pos"])["rate"].mean()
print("p0", list(means.loc["p0"].values))
print("p1", list(means.loc["p1"].values))
print("mean overall")
print(error_rates.groupby("error-rate")["rate"].mean())

import pandas as pd
import matplotlib as mpl
mpl.use("agg")
import matplotlib.pyplot as plt
import seaborn as sns

sns.set(style="ticks", palette="colorblind", context=snakemake.wildcards.context)

codebook = pd.read_table(snakemake.input.codebook, index_col=0)

all_known_counts = []
raw_errors = []
for raw_counts, known_counts in zip(snakemake.input.raw_counts, snakemake.input.known_counts):
    known_counts = pd.read_table(known_counts, index_col=[0, 1])
    known_counts = known_counts.reindex(codebook.index, level=1)
    all_known_counts.append(known_counts)

    raw_counts = pd.read_table(raw_counts, index_col=[0, 1])
    raw_counts = raw_counts["exact"] + raw_counts["corrected"]
    raw_counts = raw_counts.reindex(known_counts.index, fill_value=0)

    raw_errors.append(raw_counts - known_counts["count"])
raw_error_mean = pd.concat(raw_errors).mean()

errors = []

for uncertainty in [0, 5, 10, 20, 30]:
    u = "err-{}%".format(uncertainty) if uncertainty > 0 else "default"
    for mean, posterior_counts, known_counts in zip(snakemake.params.means, snakemake.input.get(u), all_known_counts):
        posterior_estimates = pd.read_table(posterior_counts, index_col=[0, 1])
        posterior_estimates = posterior_estimates.reindex(known_counts.index, fill_value=0)

        errors.append(pd.DataFrame({"error": posterior_estimates["expr_map"] - known_counts["count"], "mean": mean, "uncertainty": uncertainty}))

errors = pd.concat(errors)


x, y = snakemake.config["plots"]["figsize"]
plt.figure(figsize=(x * 1.5, y))
colors = sns.xkcd_palette(["light red"])
sns.violinplot(x="uncertainty", y="error", data=errors, bw=1, inner="quartile", palette=colors, linewidth=1)
plt.plot(plt.xlim(), [0, 0], "-k", linewidth=1, zorder=-5)
plt.plot(plt.xlim(), [raw_error_mean] * 2, ":k", linewidth=1, zorder=-5)
sns.despine()

plt.xlabel("error rate underestimation (%)")
plt.ylabel("predicted - truth")

plt.savefig(snakemake.output[0], bbox_inches="tight")

import matplotlib
matplotlib.use("agg")
import matplotlib.pyplot as plt
import seaborn as sns
import pandas as pd


sns.set(style="ticks", palette="colorblind", context="paper")
plt.figure(figsize=snakemake.config["plots"]["figsize"])


for f in snakemake.input:
    counts = pd.read_table(f, index_col=[0, 1])
    plt.scatter(counts["exact"], counts["corrected"], s=2, c="k", alpha=0.6, edgecolors="face", rasterized=True)

plt.xlabel("exact counts")
plt.ylabel("corrected counts")
plt.xlim((0, 1000))
plt.ylim((0, 1000))
plt.locator_params(nbins=4)
#ax = plt.gca()
#ax.set_yscale('log')
#ax.set_xscale('log')
sns.despine()
plt.savefig(snakemake.output[0], bbox_inches="tight")

import numpy as np
import pandas as pd
import matplotlib
matplotlib.use("agg")
import matplotlib.pyplot as plt
import seaborn as sns


sns.set(style="ticks", palette="colorblind", context=snakemake.wildcards.context)
plt.figure(figsize=snakemake.config["plots"]["figsize"])

means = []
stds = []
for f in snakemake.input:
    exprs = pd.read_table(f, index_col=0)
    for cell in exprs:
        expr = exprs[cell]
        means.append(expr.mean())
        stds.append(expr.std())
        expr = np.log10(expr + 1)
        sns.kdeplot(expr, color="black", clip=[0, expr.max()], alpha=0.1, lw=1, label="", kernel="gau", bw="scott")

means = np.array(means)
stds = np.array(stds)
cvs = stds / means
sns.kdeplot(np.log10(means + 1), color="red", clip=[0, means.max()], label="means")
sns.kdeplot(np.log10(stds + 1), color="red", linestyle="--", clip=[0, stds.max()], label="standard deviations")


m = means.mean()
cv = cvs.mean()
print(m, cv)
#def nbinom_params(mean, cv):
#    variance = (cv * mean) ** 2
#    p = 1 - ((variance - mean) / variance)
#    n = (mean * p) / (1 - p)
#    return n, p

#d = np.random.negative_binomial(*nbinom_params(m, cv), 5000)
#print(np.mean(d), np.std(d), np.std(d) / np.mean(d))
#sns.kdeplot(np.log10(d + 1), color="red", linestyle=":", clip=[0, d.max()], label="simulated example")


plt.xlim([0, plt.xlim()[1]])
plt.xlabel("log10 counts")
plt.ylabel("density")
plt.legend(loc="upper right")
sns.despine()

plt.savefig(snakemake.output[0], bbox_inches="tight")

import matplotlib
matplotlib.use("agg")
import matplotlib.pyplot as plt
import seaborn as sns
import pandas as pd
idx = pd.IndexSlice

import merfishtools


sns.set(style="ticks", palette="colorblind", context=snakemake.wildcards.context)
plt.figure(figsize=snakemake.config["plots"]["figsize"])

cdf = merfishtools.read_cdf(snakemake.input.fc).loc[snakemake.wildcards.gene]
est = merfishtools.read_diffexp_estimates(snakemake.input.fc_est).loc[snakemake.wildcards.gene]

merfishtools.plot_cdf(cdf, map_value=est["log2fc_map"], credible_interval=est[["log2fc_ci_lower", "log2fc_ci_upper"]], legend=snakemake.wildcards.legend == "legend")

plt.xlabel("log2 fold change")
plt.ylabel("CDF")
if snakemake.wildcards.legend == "legend":
    plt.legend(loc="best")

sns.despine()
plt.savefig(snakemake.output[0], bbox_inches="tight")