"""

gas
===

Functions for plotting gas quantities

"""

import logging

import matplotlib.pyplot as plt
import numpy as np

from ..analysis import angmom, halo, profile
from ..units import Unit
from .generic import hist2d

logger = logging.getLogger('pynbody.plot.gas')


def rho_T(sim, rho_units=None, rho_range=None, t_range=None, two_phase='split', **kwargs):
    """
    Plot Temperature vs. Density for the gas particles in the snapshot.

    **Optional keywords:**

       *rho_units*: specify the density units (default is the same units as the current 'rho' array)

       *t_range*: list, array, or tuple
          ``size(t_range)`` must be 2. Specifies the temperature range.

       *rho_range:* tuple
          ``size(rho_range)`` must be 2. Specifies the density range.

       *two_phase*: if two-phase particles are detected, either plot each phase separately ('split'), or merge them ('merge')

    See :func:`~pynbody.plot.generic.hist2d` for other plotting keyword options


    """
    if rho_units is None:
        rho_units = sim.gas['rho'].units

    if t_range is not None:
        kwargs['y_range'] = t_range
        assert len(kwargs['y_range']) == 2

    if rho_range is not None:
        kwargs['x_range'] = rho_range
        assert len(kwargs['x_range']) == 2
    else:
        rho_range = False

    if 'xlabel' in kwargs:
        xlabel = kwargs['xlabel']
        del kwargs['xlabel']
    else:
        xlabel = r'log$_{10}$($\rho$/$' + Unit(rho_units).latex() + '$)'

    if 'ylabel' in kwargs:
        ylabel = kwargs['ylabel']
        del kwargs['ylabel']
    else:
        ylabel = r'log$_{10}$(T/$' + sim.gas['temp'].units.latex() + '$)'

    if 'Tinc' in sim.loadable_keys() and two_phase == 'merge':
        return hist2d(sim.gas['rho'].in_units(rho_units),sim.gas['Tinc'],
                      xlogrange=True,ylogrange=True,xlabel=xlabel,
                      ylabel=ylabel, **kwargs)

    if 'uHot' in sim.loadable_keys() and 'MassHot' in sim.loadable_keys() and two_phase == 'split':
        E = sim.g['uHot']*sim.g['MassHot']+sim.g['u']*(sim.g['mass']-sim.g['MassHot'])
        rho = np.concatenate((sim.g['rho'].in_units(rho_units)*E/(sim.g['mass']*sim.g['u']),
            sim.g['rho'].in_units(rho_units)*E/(sim.g['mass']*sim.g['uHot'])))
        temp = np.concatenate((sim.g['temp'], sim.g['temp']/sim.g['u']*sim.g['uHot']))
        temp = temp[np.where(np.isfinite(rho))]
        rho = rho[np.where(np.isfinite(rho))]
        return hist2d(rho, temp, xlogrange=True,ylogrange=True,xlabel=xlabel,
                ylabel=ylabel, **kwargs)

    return hist2d(sim.gas['rho'].in_units(rho_units),sim.gas['temp'],
                  xlogrange=True,ylogrange=True,xlabel=xlabel,
                  ylabel=ylabel, **kwargs)



def temp_profile(sim, center=True, r_units='kpc', bin_spacing='equaln',
                 clear=True, filename=None, **kwargs):
    """

    Centre on potential minimum, align so that the disk is in the
    x-y plane, then plot the temperature profile as a
    function of radius.

    """

    if center:
        angmom.sideon(sim)

    if 'min' in kwargs:
        min_r = kwargs['min']
    else:
        min_r = sim['r'].min()
    if 'max' in kwargs:
        max_r = kwargs['max']
    else:
        max_r = sim['r'].max()

    pro = profile.Profile(sim.gas, type=bin_spacing, rmin =min_r, rmax =max_r)

    r = pro['rbins'].in_units(r_units)
    tempprof = pro['temp']

    if clear:
        plt.clf()

    plt.semilogy(r, tempprof)

    plt.xlabel("r / $" + r.units.latex() + "$")
    plt.ylabel("Temperature [K]")
    if (filename):
        logger.info("Saving %s", filename)
        plt.savefig(filename)
