import multiprocessing
import numpy as np

from .model import Model
from ..grid import SphericalPolarGrid
from ..util.interpolate import interp1d_fast
from .analytical_yso_model import AnalyticalYSOModel


def find_last_iteration(file_handle):
    max_iteration = 0
    for group_name in file_handle:
        if "iteration" in group_name:
            iteration = int(group_name.split('_')[1])
            max_iteration = max(iteration, max_iteration)
    return max_iteration


def tau_to_radius(model, tau, wav):
    """
    Given a Model instance with a spherical polar coordinate grid, find the
    radius from which the optical depth to escape radially is a fixed value.

    This only works for spherical polar grids, but works for 1-, 2-, and
    3-d grids.

    Parameters
    ----------
    model : `~hyperion.model.Model` instance
    tau : float
        The optical depth for which to find the surface
    wav : float
        The wavelength at which the optical depth is defined

    Returns
    -------
    r : np.ndarray
        The radius or radii at which the optical depth to escape radially
        is ``tau`` at ``wav``. This is a scalar, a 1-d, or a 2-d array
        depending on the dimensionality of the grid.
    """

    if not isinstance(model, Model):
        raise TypeError("model should be a Model instance")

    if model.grid is None:
        raise Exception("Coordinate grid has not been defined")

    if not isinstance(model.grid, SphericalPolarGrid):
        raise TypeError("This method can only be called for spherical polar grids")

    # Initialize cumulative optical depth array
    tau_all = np.zeros(model.grid.shape)

    # Loop over envelopes and add cumulative column density
    for i, item in enumerate(model.grid['density']):

        # Find density
        rho = item.array

        # Find optical depth in all cells in radial direction
        dtau = model.grid.widths[0, :, :, :] * rho * model.dust[i].optical_properties.interp_chi_wav(wav)

        # Find cumulative sum starting from the ouside
        tau_all += np.cumsum(dtau[:, :, ::-1], axis=2)

    r = np.zeros(model.grid.shape[:2])
    for ip in range(model.grid.shape[0]):