from __future__ import absolute_import
import os.path as op
import shutil

import numpy as nm

from sfepy.base.base import ordered_iteritems, output, set_defaults, assert_
from sfepy.base.base import Struct
from sfepy.homogenization.engine import HomogenizationEngine
from sfepy.homogenization.homogen_app import HomogenizationApp
from sfepy.homogenization.coefficients import Coefficients
from sfepy.homogenization.coefs_base import CoefDummy
from sfepy.applications import PDESolverApp
from sfepy.base.plotutils import plt
import six
from six.moves import range

def try_set_defaults(obj, attr, defaults, recur=False):
    try:
        values = getattr(obj, attr)

    except:
        values = defaults

    else:
        if recur and isinstance(values, dict):
            for key, val in six.iteritems(values):
                set_defaults(val, defaults)

        else:
            set_defaults(values, defaults)

    return values

def save_raw_bg_logs(filename, logs):
    """
    Save raw band gaps `logs` into the `filename` file.
    """
    out = {}

    iranges = nm.cumsum([0] + [len(ii) for ii in logs.freqs])
    out['iranges'] = iranges
    for key, log in ordered_iteritems(logs.to_dict()):
        out[key] = nm.concatenate(log, axis=0)

    with open(filename, 'w') as fd:
        nm.savez(fd, **out)

def transform_plot_data(datas, plot_transform, conf):
    if plot_transform is not None:
        fun = conf.get_function(plot_transform[0])

    dmin, dmax = 1e+10, -1e+10
    tdatas = []
    for data in datas:
        tdata = data.copy()
        if plot_transform is not None:
            tdata = fun(tdata, *plot_transform[1:])
        dmin = min(dmin, nm.nanmin(tdata))
        dmax = max(dmax, nm.nanmax(tdata))
        tdatas.append(tdata)
    dmin, dmax = min(dmax - 1e-8, dmin), max(dmin + 1e-8, dmax)
    return (dmin, dmax), tdatas

def plot_eigs(fig_num, plot_rsc, plot_labels, valid, freq_range, plot_range,
              show=False, clear=False, new_axes=False):
    """
    Plot resonance/eigen-frequencies.

    `valid` must correspond to `freq_range`

    resonances : red
    masked resonances: dotted red
    """
    if plt is None: return
    assert_(len(valid) == len(freq_range))

    fig = plt.figure(fig_num)
    if clear:
        fig.clf()
    if new_axes:
        ax = fig.add_subplot(111)
    else:
        ax = fig.gca()

    l0 = l1 = None
    for ii, f in enumerate(freq_range):
        if valid[ii]:
            l0 = ax.plot([f, f], plot_range, **plot_rsc['resonance'])[0]
        else:
            l1 = ax.plot([f, f], plot_range, **plot_rsc['masked'])[0]

    if l0:
        l0.set_label(plot_labels['resonance'])
    if l1:
        l1.set_label(plot_labels['masked'])

    if new_axes:
        ax.set_xlim([freq_range[0], freq_range[-1]])
        ax.set_ylim(plot_range)

    if show:
        plt.show()
    return fig

def plot_logs(fig_num, plot_rsc, plot_labels,
              freqs, logs, valid, freq_range, plot_range,
              draw_eigs=True, show_legend=True, show=False,
              clear=False, new_axes=False):
    """
    Plot logs of min/middle/max eigs of a mass matrix.
    """
    if plt is None: return

    fig = plt.figure(fig_num)
    if clear:
        fig.clf()
    if new_axes:
        ax = fig.add_subplot(111)
    else:
        ax = fig.gca()

    if draw_eigs:
        plot_eigs(fig_num, plot_rsc, plot_labels, valid, freq_range,
                  plot_range)

    for ii, log in enumerate(logs):
        l1 = ax.plot(freqs[ii], log[:, -1], **plot_rsc['eig_max'])

        if log.shape[1] >= 2:
            l2 = ax.plot(freqs[ii], log[:, 0], **plot_rsc['eig_min'])
        else:
            l2 = None

        if log.shape[1] == 3:
            l3 = ax.plot(freqs[ii], log[:, 1], **plot_rsc['eig_mid'])
        else:
            l3 = None

    l1[0].set_label(plot_labels['eig_max'])
    if l2:
        l2[0].set_label(plot_labels['eig_min'])
    if l3:
        l3[0].set_label(plot_labels['eig_mid'])

    fmin, fmax = freqs[0][0], freqs[-1][-1]
    ax.plot([fmin, fmax], [0, 0], **plot_rsc['x_axis'])

    ax.set_xlabel(plot_labels['x_axis'])
    ax.set_ylabel(plot_labels['y_axis'])

    if new_axes:
        ax.set_xlim([fmin, fmax])
        ax.set_ylim(plot_range)

    if show_legend:
        ax.legend()

    if show:
        plt.show()
    return fig

def plot_gap(ax, ranges, kind, kind_desc, plot_range, plot_rsc):
    """
    Plot single band gap frequency ranges as rectangles.
    """
    def draw_rect(ax, x, y, rsc):
        ax.fill(nm.asarray(x)[[0,1,1,0]],
                 nm.asarray(y)[[0,0,1,1]],
                 **rsc)

    # Colors.
    strong = plot_rsc['strong_gap']
    weak = plot_rsc['weak_gap']
    propagation = plot_rsc['propagation']

    if kind == 'p':
        draw_rect(ax, ranges[0], plot_range, propagation)
    elif kind == 'w':
        draw_rect(ax, ranges[0], plot_range, weak)
    elif kind == 'wp':
        draw_rect(ax, ranges[0], plot_range, weak)
        draw_rect(ax, ranges[1], plot_range, propagation)
    elif kind == 's':
        draw_rect(ax, ranges[0], plot_range, strong)
    elif kind == 'sw':
        draw_rect(ax, ranges[0], plot_range, strong)
        draw_rect(ax, ranges[1], plot_range, weak)
    elif kind == 'swp':
        draw_rect(ax, ranges[0], plot_range, strong)
        draw_rect(ax, ranges[1], plot_range, weak)
        draw_rect(ax, ranges[2], plot_range, propagation)
    elif kind == 'is':
        draw_rect(ax, ranges[0], plot_range, strong)
    elif kind == 'iw':
        draw_rect(ax, ranges[0], plot_range, weak)
    else:
        raise ValueError('unknown band gap kind! (%s)' % kind)

def plot_gaps(fig_num, plot_rsc, gaps, kinds, gap_ranges, freq_range,
              plot_range, show=False, clear=False, new_axes=False):
    """
    Plot band gaps as rectangles.
    """
    if plt is None: return

    fig = plt.figure(fig_num)
    if clear:
        fig.clf()
    if new_axes:
        ax = fig.add_subplot(111)
    else:
        ax = fig.gca()

    for ii in range(len(freq_range) - 1):
        f0, f1 = freq_range[[ii, ii+1]]
        gap = gaps[ii]
        ranges = gap_ranges[ii]

        if isinstance(gap, list):
            for ig, (gmin, gmax) in enumerate(gap):
                kind, kind_desc = kinds[ii][ig]
                plot_gap(ax, ranges[ig], kind, kind_desc, plot_range, plot_rsc)

                output(ii, gmin[0], gmax[0], '%.8f' % f0, '%.8f' % f1)
                output(' -> %s\n    %s' %(kind_desc, ranges[ig]))
        else:
            gmin, gmax = gap
            kind, kind_desc = kinds[ii]
            plot_gap(ax, ranges, kind, kind_desc, plot_range, plot_rsc)
            output(ii, gmin[0], gmax[0], '%.8f' % f0, '%.8f' % f1)
            output(' -> %s\n    %s' %(kind_desc, ranges))

    if new_axes:
        ax.set_xlim([freq_range[0], freq_range[-1]])
        ax.set_ylim(plot_range)

    if show:
        
plt.show()