#!/usr/bin/env python

"""Tests for pyani graphics

If the test is run directly at the command-line, the output obtained by each
test is returned to STDOUT.
"""

import os
import pandas as pd

from pyani import pyani_graphics, pyani_config, pyani_tools


# Work out where we are. We need to do this to find related data files
# for testing
curdir = os.path.dirname(os.path.abspath(__file__))

OUTDIR = os.path.join("tests", "test_graphics_output")


def define_inputs():
    return {
        "infilename": os.path.join("tests", "target_ANIm_output", "ANIm_percentage_identity.tab"),
        "labels": pyani_tools.get_labels(os.path.join("tests", "test_ani_data", "labels.tab"), logger=None),
        "classes": pyani_tools.get_labels(os.path.join("tests", "test_ani_data", "classes.tab"), logger=None),
    }


def draw_format_method(fmt, mth):
    """Render graphics format and method output."""
    inputs = define_inputs()
    outfilename = os.path.join(OUTDIR, "%s.%s" % (mth, fmt))
    stem = "ANIm_percentage_identity"
    df = pd.read_csv(inputs["infilename"], index_col=0, sep="\t")
    os.makedirs(OUTDIR, exist_ok=True)
    fn = {"mpl": pyani_graphics.heatmap_mpl, "seaborn": pyani_graphics.heatmap_seaborn}
    params = {"mpl": pyani_config.params_mpl, "seaborn": pyani_config.params_mpl}
    method_params = pyani_graphics.Params(params[mth](df)[stem], inputs["labels"], inputs["classes"])
    fn[mth](df, outfilename, title="%s:%s test" % (mth, fmt), params=method_params)


def test_png_mpl():
    """Write .png graphics with mpl"""
    draw_format_method("png", "mpl")


def test_svg_mpl():
    """Write .svg graphics with mpl"""
    draw_format_method("svg", "mpl")


def test_pdf_mpl():
    """Write .pdf graphics with mpl"""
    draw_format_method("pdf", "mpl")


def test_png_seaborn():
    """Write .png graphics with seaborn"""
    draw_format_method("png", "seaborn")


def test_svg_seaborn():
    """Write .svg graphics with seaborn"""
    draw_format_method("svg", "seaborn")


def test_pdf_seaborn():
    """Write .pdf graphics with seaborn"""
    draw_format_method("pdf", "seaborn")