import functools
import math
import numbers
import os
import re
import sys
import time

import h5py
import numpy as np
from meep.geom import init_do_averaging
from meep.simulation import get_num_args
from meep.verbosity_mgr import Verbosity

import meep as mp

from . import mode_solver, with_hermitian_epsilon

try:
    basestring
except NameError:
    basestring = str

U_MIN = 0
U_PROD = 1
U_MEAN = 2

verbosity = Verbosity(mp.cvar, "meep", 1)


class MPBArray(np.ndarray):
    def __new__(cls, input_array, lattice, kpoint=None, bloch_phase=False):
        # Input array is an already formed ndarray instance
        # We first cast to be our class type
        obj = np.asarray(input_array).view(cls)
        # add the new properties to the created instance
        obj.lattice = lattice
        obj.kpoint = kpoint
        obj.bloch_phase = bloch_phase
        # Finally, we must return the newly created object:
        return obj

    def __array_finalize__(self, obj):
        # ``self`` is a new object resulting from
        # ndarray.__new__(MPBArray, ...), therefore it only has
        # attributes that the ndarray.__new__ constructor gave it -
        # i.e. those of a standard ndarray.

        # We could have got to the ndarray.__new__ call in 3 ways:
        # From an explicit constructor - e.g. MPBArray(lattice):
        #    obj is None
        #    (we're in the middle of the MPBArray.__new__
        #    constructor, and self.lattice will be set when we return to
        #    MPBArray.__new__)
        if obj is None:
            return

        # From view casting - e.g arr.view(MPBArray):
        #    obj is arr
        #    (type(obj) can be MPBArray)
        # From new-from-template - e.g mpbarr[:3]
        #    type(obj) is MPBArray
        #
        # Note that it is here, rather than in the __new__ method,
        # that we set the default value for 'lattice', because this
        # method sees all creation of default objects - with the
        # MPBArray.__new__ constructor, but also with
        # arr.view(MPBArray).
        self.lattice = getattr(obj, "lattice", None)
        self.kpoint = getattr(obj, "kpoint", None)
        self.bloch_phase = getattr(obj, "bloch_phase", False)


class ModeSolver:
    def __init__(
        self,
        resolution=10,
        is_negative_epsilon_ok=False,
        eigensolver_flops=0,
        eigensolver_flags=68,
        use_simple_preconditioner=False,
        force_mu=False,
        mu_input_file="",
        epsilon_input_file="",
        mesh_size=3,
        target_freq=0.0,
        tolerance=1.0e-7,
        num_bands=1,
        k_points=None,
        ensure_periodicity=True,
        geometry=None,
        geometry_lattice=mp.Lattice(),
        geometry_center=mp.Vector3(0, 0, 0),
        default_material=mp.Medium(epsilon=1),
        dimensions=3,
        random_fields=False,
        filename_prefix="",
        deterministic=False,
        verbose=False,
        optimize_grid_size=True,
        eigensolver_nwork=3,
        eigensolver_block_size=-11,
    ):

        self.mode_solver = None
        self.resolution = resolution
        self.eigensolver_flags = eigensolver_flags
        self.k_points = k_points or []
        self.geometry = geometry or []
        self.geometry_lattice = geometry_lattice
        self.geometry_center = mp.Vector3(*geometry_center)
        self.default_material = default_material
        self.random_fields = random_fields
        self.filename_prefix = filename_prefix
        self.optimize_grid_size = optimize_grid_size
        self.parity = ""
        self.iterations = 0
        self.all_freqs = None
        self.freqs = []
        self.band_range_data = []
        self.total_run_time = 0
        self.current_k = mp.Vector3()
        self.k_split_num = 1
        self.k_split_index = 0
        self.eigensolver_iters = []

        grid_size = self._adjust_grid_size()

        if type(self.default_material) is not mp.Medium and callable(
            self.default_material
        ):
            init_do_averaging(self.default_material)
            self.default_material.eps = False

        self.mode_solver = mode_solver(
            num_bands,
            self.resolution,
            self.geometry_lattice,
            tolerance,
            mesh_size,
            self.default_material,
            deterministic,
            target_freq,
            dimensions,
            verbose,
            ensure_periodicity,
            eigensolver_flops,
            is_negative_epsilon_ok,
            epsilon_input_file,
            mu_input_file,
            force_mu,
            use_simple_preconditioner,
            grid_size,
            eigensolver_nwork,
            eigensolver_block_size,
        )

    @property
    def num_bands(self):
        return self.mode_solver.num_bands

    @num_bands.setter
    def num_bands(self, val):
        self.mode_solver.set_num_bands(val)

    @property
    def resolution(self):
        return self._resolution

    @resolution.setter
    def resolution(self, val):
        if isinstance(val, numbers.Number):
            self._resolution = [val, val, val]
        elif isinstance(val, mp.Vector3):
            self._resolution = [val.x, val.y, val.z]
        else:
            t = type(val)
            raise TypeError(f"resolution must be a number or a Vector3: Got {t}")

        if self.mode_solver:
            self.mode_solver.resolution = self._resolution
            grid_size = self._adjust_grid_size()
            self.mode_solver.set_grid_size(grid_size)

    @property
    def geometry_lattice(self):
        return self._geometry_lattice

    @geometry_lattice.setter
    def geometry_lattice(self, val):
        self._geometry_lattice = val
        if self.mode_solver:
            self.mode_solver.set_libctl_geometry_lattice(val)
            grid_size = self._adjust_grid_size()
            self.mode_solver.set_grid_size(grid_size)

    @property
    def tolerance(self):
        return self.mode_solver.tolerance

    @tolerance.setter
    def tolerance(self, val):
        self.mode_solver.tolerance = val

    @property
    def mesh_size(self):
        return self.mode_solver.mesh_size

    @mesh_size.setter
    def mesh_size(self, val):
        self.mode_solver.mesh_size = val

    @property
    def deterministic(self):
        return self.mode_solver.deterministic

    @deterministic.setter
    def deterministic(self, val):
        self.mode_solver.deterministic = val

    @property
    def target_freq(self):
        return self.mode_solver.target_freq

    @target_freq.setter
    def target_freq(self, val):
        self.mode_solver.target_freq = val

    @property
    def dimensions(self):
        return self.mode_solver.get_libctl_dimensions()

    @dimensions.setter
    def dimensions(self, val):
        self.mode_solver.set_libctl_dimensions(val)

    @property
    def verbose(self):
        return self.mode_solver.verbose

    @verbose.setter
    def verbose(self, val):
        self.mode_solver.verbose = val

    @property
    def ensure_periodicity(self):
        return self.mode_solver.get_libctl_ensure_periodicity()

    @ensure_periodicity.setter
    def ensure_periodicity(self, val):
        self.mode_solver.set_libctl_ensure_periodicity(val)

    @property
    def eigensolver_flops(self):
        return self.mode_solver.eigensolver_flops

    @eigensolver_flops.setter
    def eigensolver_flops(self, val):
        self.mode_solver.eigensolver_flops = val

    @property
    def is_negative_epsilon_ok(self):
        return self.mode_solver.negative_epsilon_ok

    @is_negative_epsilon_ok.setter
    def is_negative_epsilon_ok(self, val):
        self.mode_solver.negative_epsilon_ok = val

    @property
    def epsilon_input_file(self):
        return self.mode_solver.epsilon_input_file

    @epsilon_input_file.setter
    def epsilon_input_file(self, val):
        self.mode_solver.epsilon_input_file = val

    @property
    def mu_input_file(self):
        return self.mode_solver.mu_input_file

    @mu_input_file.setter
    def mu_input_file(self, val):
        self.mode_solver.mu_input_file = val

    @property
    def force_mu(self):
        return self.mode_solver.force_mu

    @force_mu.setter
    def force_mu(self, val):
        self.mode_solver.force_mu = val

    @property
    def use_simple_preconditioner(self):
        return self.mode_solver.use_simple_preconditioner

    @use_simple_preconditioner.setter
    def use_simple_preconditioner(self, val):
        self.mode_solver.use_simple_preconditioner = val

    @property
    def eigensolver_nwork(self):
        return self.mode_solver.eigensolver_nwork

    @eigensolver_nwork.setter
    def eigensolver_nwork(self, val):
        self.mode_solver.eigensolver_nwork = val

    @property
    def eigensolver_block_size(self):
        return self.mode_solver.eigensolver_block_size

    @eigensolver_block_size.setter
    def eigensolver_block_size(self, val):
        self.mode_solver.eigensolver_block_size = val

    def _adjust_grid_size(self):
        grid_size = self._get_grid_size()

        if self.optimize_grid_size:
            grid_size = self._optimize_grid_size(grid_size)

        return grid_size

    def allow_negative_epsilon(self):
        self.is_negative_epsilon_ok = True
        self.target_freq = 1 / mp.inf

    def get_filename_prefix(self):
        if self.filename_prefix:
            return f"{self.filename_prefix}-"
        _, filename = os.path.split(sys.argv[0])

        return (
            ""
            if filename in ["ipykernel_launcher.py", "__main__.py"]
            else re.sub(r"\.py$", "", filename) + "-"
        )

    def get_freqs(self):
        return self.mode_solver.get_freqs()

    def multiply_bloch_phase(self, arr):
        dims = arr.shape
        arr = arr.ravel()
        self.mode_solver.multiply_bloch_phase(arr)
        return np.reshape(arr, dims)

    def get_poynting(self, which_band):
        e = self.get_efield(which_band, False)
        dims = e.shape
        e = e.ravel()
        h = self.get_hfield(which_band, False).ravel()
        # Reshape into rows of vector3s
        e = e.reshape((int(e.shape[0] / 3), 3))
        h = h.reshape((int(h.shape[0] / 3), 3))

        res = np.zeros(e.shape, dtype=np.complex128)

        def ExH(e, h):
            ev = mp.Vector3(e[0], e[1], e[2])
            hv = mp.Vector3(h[0], h[1], h[2])
            return ev.conj().cross(hv)

        for i in range(e.shape[0]):
            res[i] = np.array(ExH(e[i], h[i]))

        flat_res = res.ravel()
        self.mode_solver.set_curfield_cmplx(flat_res)
        self.mode_solver.set_curfield_type("v")

        arr = np.reshape(res, dims)
        return MPBArray(arr, self.get_lattice(), self.current_k)

    def get_epsilon(self):
        self.mode_solver.get_epsilon()
        return self.get_curfield_as_array(False)

    def get_mu(self):
        self.mode_solver.get_mu()
        return self.get_curfield_as_array(False)

    def get_bfield(self, which_band, bloch_phase=True):
        return self._get_field("b", which_band, bloch_phase)

    def get_efield(self, which_band, bloch_phase=True):
        return self._get_field("e", which_band, bloch_phase)

    def get_dfield(self, which_band, bloch_phase=True):
        return self._get_field("d", which_band, bloch_phase)

    def get_hfield(self, which_band, bloch_phase=True):
        return self._get_field("h", which_band, bloch_phase)

    def get_charge_density(self, which_band, bloch_phase=True):
        self.get_efield(which_band, bloch_phase)
        self.mode_solver.compute_field_divergence()

    def _get_field(self, f, band, bloch_phase):
        if self.mode_solver is None:
            raise ValueError(
                "Must call a run function before attempting to get a field"
            )

        if f == "b":
            self.mode_solver.get_bfield(band)
        elif f == "d":
            self.mode_solver.get_dfield(band)
        elif f == "e":
            self.mode_solver.get_efield(band)
        elif f == "h":
            self.mode_solver.get_hfield(band)

        dims = self.mode_solver.get_dims()

        while len(dims) < 3:
            dims += [1]

        dims += [3]
        arr = np.zeros(np.prod(dims), np.complex128)

        if bloch_phase:
            self.mode_solver.multiply_bloch_phase()

        self.mode_solver.get_curfield_cmplx(arr)

        arr = np.reshape(arr, dims)
        return MPBArray(
            arr, self.get_lattice(), self.current_k, bloch_phase=bloch_phase
        )

    def get_curfield_as_array(self, bloch_phase=True):
        dims = self.mode_solver.get_dims()
        arr = np.zeros(np.prod(dims))
        self.mode_solver.get_curfield(arr)
        arr = np.reshape(arr, dims)

        return MPBArray(
            arr, self.get_lattice(), self.current_k, bloch_phase=bloch_phase
        )

    def get_dpwr(self, band):
        self.get_dfield(band, False)
        self.compute_field_energy()
        return self.get_curfield_as_array(False)

    def get_bpwr(self, band):
        self.get_bfield(band, False)
        self.compute_field_energy()
        return self.get_curfield_as_array(False)

    def fix_field_phase(self):
        self.mode_solver.fix_field_phase()

    def get_epsilon_point(self, p):
        return self.mode_solver.get_epsilon_point(p)

    def get_epsilon_inverse_tensor_point(self, p):
        return self.mode_solver.get_epsilon_inverse_tensor_point(p)

    def get_energy_point(self, p):
        return self.mode_solver.get_energy_point(p)

    def get_field_point(self, p):
        return self.mode_solver.get_field_point(p)

    def get_bloch_field_point(self, p):
        return self.mode_solver.get_bloch_field_point(p)

    def get_tot_pwr(self, which_band):
        epwr = self.get_dpwr(which_band)
        hpwr = self.get_bpwr(which_band)

        tot_pwr = epwr + hpwr

        self.mode_solver.set_curfield(tot_pwr.ravel())
        self.mode_solver.set_curfield_type("R")

        return MPBArray(tot_pwr, self.get_lattice(), self.current_k, bloch_phase=False)

    def get_eigenvectors(self, first_band, num_bands):
        dims = self.mode_solver.get_eigenvectors_slice_dims(num_bands)
        ev = np.zeros(np.prod(dims), dtype=np.complex128)
        self.mode_solver.get_eigenvectors(first_band - 1, num_bands, ev)
        return MPBArray(ev.reshape(dims), self.get_lattice(), self.current_k)

    def set_eigenvectors(self, ev, first_band):
        self.mode_solver.set_eigenvectors(first_band - 1, ev.flatten())

    def save_eigenvectors(self, filename):
        with h5py.File(filename, "w") as f:
            ev = self.get_eigenvectors(1, self.num_bands)
            f["rawdata"] = ev

    def load_eigenvectors(self, filename):
        with h5py.File(filename, "r") as f:
            ev = f["rawdata"][()]
            self.set_eigenvectors(ev, 1)
            self.mode_solver.curfield_reset()

    # The band-range-data is a list of tuples, each consisting of a (min, k-point)
    # tuple and a (max, k-point) tuple, with each min/max pair describing the
    # frequency range of a band and the k-points where it achieves its minimum/maximum.
    # Here, we update this data with a new list of band frequencies, and return the new
    # data.  If band-range-data is null or too short, the needed entries will be created.
    def update_band_range_data(self, brd, freqs, kpoint):
        def update_brd(brd, freqs, br_start):
            if not freqs:
                return br_start + brd
            br = brd[0] if brd else ((mp.inf, -1), (-mp.inf, -1))
            br_rest = brd[1:] if brd else []
            newmin = (freqs[0], kpoint) if freqs[0] < br[0][0] else br[0]
            newmax = (freqs[0], kpoint) if freqs[0] > br[1][0] else br[1]
            new_start = br_start + [(newmin, newmax)]
            return update_brd(br_rest, freqs[1:], new_start)

        return update_brd(brd, freqs, [])

    def output_band_range_data(self, br_data):
        if verbosity.mpb > 0:
            fmt = "Band {} range: {} at {} to {} at {}"
            for tup, band in zip(br_data, range(1, len(br_data) + 1)):
                min_band, max_band = tup
                min_freq, min_kpoint = min_band
                max_freq, max_kpoint = max_band
                print(fmt.format(band, min_freq, min_kpoint, max_freq, max_kpoint))

    # Output any gaps in the given band ranges, and return a list of the gaps as
    # a list of (percent, freq-min, freq-max) tuples.
    def output_gaps(self, br_data):
        def ogaps(br_cur, br_rest, i, gaps):
            if not br_rest:
                gaps = list(reversed(gaps))
                return [
                    (gaps[i + 2], gaps[i + 1], gaps[i]) for i in range(0, len(gaps), 3)
                ]
            else:
                br_rest_min_f = br_rest[0][0][0]
                br_cur_max_f = br_cur[1][0]
                if br_cur_max_f >= br_rest_min_f:
                    return ogaps(br_rest[0], br_rest[1:], i + 1, gaps)
                gap_size = (200 * (br_rest_min_f - br_cur_max_f)) / (
                    br_rest_min_f + br_cur_max_f
                )
                if verbosity.mpb > 0:
                    fmt = "Gap from band {} ({}) to band {} ({}), {}%"
                    print(fmt.format(i, br_cur_max_f, i + 1, br_rest_min_f, gap_size))
                return ogaps(
                    br_rest[0],
                    br_rest[1:],
                    i + 1,
                    [gap_size, br_cur_max_f, br_rest_min_f] + gaps,
                )

        if not br_data:
            return []
        else:
            return ogaps(br_data[0], br_data[1:], 1, [])

    # Return the frequency gap from the band #lower-band to the band
    # #(lower-band+1), as a percentage of mid-gap frequency.  The "gap"
    # may be negative if the maximum of the lower band is higher than the
    # minimum of the upper band.  (The gap is computed from the
    # band-range-data of the previous run.)
    def retrieve_gap(self, lower_band):
        if lower_band + 1 > len(self.band_range_data):
            raise ValueError("retrieve-gap called for higher band than was calculated")

        f1 = self.band_range_data[lower_band - 1][1][0]
        f2 = self.band_range_data[lower_band][0][0]

        return (f2 - f1) / (0.005 * (f1 + f2))

    # Split a list L into num more-or-less equal pieces, returning the piece
    # given by index (in 0..num-1), along with the index in L of the first
    # element of the piece, as a list: [first-index, piece-of-L]
    def list_split(self, l, num, index):
        def list_sub(l, start, length, index, rest):
            if not l:
                return list(reversed(rest))
            if index >= start and index < (start + length):
                return list_sub(l[1:], start, length, index + 1, [l[0]] + rest)
            else:
                return list_sub(l[1:], start, length, index + 1, rest)

        if index >= num or index < 0:
            return (len(l), [])
        else:
            block_size = (len(l) + num - 1) // num
            start = index * block_size
            length = min(block_size, (len(l) - index * block_size))
            return (start, list_sub(l, start, length, 0, []))

    def get_lattice(self):
        if self.mode_solver is None:
            raise RuntimeError("Must call ModeSolver.run before getting the lattice.")

        lattice = np.zeros((3, 3))
        self.mode_solver.get_lattice(lattice)

        return lattice

    def output_field(self):
        self.output_field_to_file(mp.ALL, self.get_filename_prefix())

    def output_field_x(self):
        self.output_field_to_file(0, self.get_filename_prefix())

    def output_field_y(self):
        self.output_field_to_file(1, self.get_filename_prefix())

    def output_field_z(self):
        self.output_field_to_file(2, self.get_filename_prefix())

    def output_epsilon(self):
        self.mode_solver.get_epsilon()
        self.output_field_to_file(mp.ALL, self.get_filename_prefix())

    def output_mu(self):
        self.mode_solver.get_mu()
        self.output_field_to_file(mp.ALL, self.get_filename_prefix())

    def output_field_to_file(self, component, fname_prefix):
        curfield_type = self.mode_solver.get_curfield_type()
        output_k = self.mode_solver.get_output_k()

        if curfield_type in "Rv":
            # Generic scalar/vector field. Don't know k
            output_k = [0, 0, 0]

        if curfield_type in "dhbecv":
            self._output_vector_field(curfield_type, fname_prefix, output_k, component)
        elif curfield_type == "C":
            self._output_complex_scalar_field(fname_prefix, output_k)
        elif curfield_type in "DHBnmR":
            self._output_scalar_field(curfield_type, fname_prefix)
        else:
            raise ValueError(f"Unkown field type: {curfield_type}")

        self.mode_solver.curfield_reset()

    def _output_complex_scalar_field(self, fname_prefix, output_k):
        curfield_type = "C"
        kpoint_index = self.mode_solver.get_kpoint_index()
        curfield_band = self.mode_solver.curfield_band
        fname = f"{curfield_type}.k{kpoint_index:02d}.b{curfield_band:02d}"
        description = "{} field, kpoint {}, band {}, freq={:.6g}".format(
            curfield_type, kpoint_index, curfield_band, self.freqs[curfield_band - 1]
        )
        fname = self._create_fname(fname, fname_prefix, True)
        if verbosity.mpb > 0:
            print(f"Outputting complex scalar field to {fname}...")

        with h5py.File(fname, "w") as f:
            f["description"] = description.encode()
            f["Bloch wavevector"] = np.array(output_k)
            self._write_lattice_vectors(f)

            dims = self.mode_solver.get_dims()
            field = np.empty(np.prod(dims), np.complex128)
            self.mode_solver.get_curfield_cmplx(field)

            reshaped_field = field.reshape(dims)

            f["c.r"] = np.real(reshaped_field)
            f["c.i"] = np.imag(reshaped_field)

    def _output_vector_field(self, curfield_type, fname_prefix, output_k, component):
        components = ["x", "y", "z"]
        kpoint_index = self.mode_solver.get_kpoint_index()
        curfield_band = self.mode_solver.curfield_band
        fname = f"{curfield_type}.k{kpoint_index:02d}.b{curfield_band:02d}"

        if component >= 0:
            fname += f".{components[component]}"

        description = "{} field, kpoint {}, band {}, freq={:.6g}".format(
            curfield_type, kpoint_index, curfield_band, self.freqs[curfield_band - 1]
        )

        fname = self._create_fname(fname, fname_prefix, True)
        if verbosity.mpb > 0:
            print(f"Outputting fields to {fname}...")

        with h5py.File(fname, "w") as f:
            f["description"] = description.encode()
            f["Bloch wavevector"] = np.array(output_k)
            self._write_lattice_vectors(f)

            if curfield_type != "v":
                self.mode_solver.multiply_bloch_phase()

            for c_idx, c in enumerate(components):
                if component >= 0 and c_idx != component:
                    continue

                dims = self.mode_solver.get_dims()
                field = np.empty(np.prod(dims) * 3, np.complex128)
                self.mode_solver.get_curfield_cmplx(field)
                component_field = field[c_idx::3].reshape(dims)

                name = f"{c}.r"
                f[name] = np.real(component_field)
                name = f"{c}.i"
                f[name] = np.imag(component_field)

    def _output_scalar_field(self, curfield_type, fname_prefix):
        components = ["x", "y", "z"]

        if curfield_type == "n":
            fname = "epsilon"
            description = "dielectric function, epsilon"
        elif curfield_type == "m":
            fname = "mu"
            description = "permeability mu"
        else:
            kpoint_index = self.mode_solver.get_kpoint_index()
            curfield_band = self.mode_solver.curfield_band
            fname = "{}pwr.k{:02d}.b{:02d}".format(
                curfield_type.lower(), kpoint_index, curfield_band
            )
            descr_fmt = "{} field energy density, kpoint {}, band {}, freq={:.6g}"
            description = descr_fmt.format(
                curfield_type,
                kpoint_index,
                curfield_band,
                self.freqs[curfield_band - 1],
            )

        parity_suffix = curfield_type not in "mn"
        fname = self._create_fname(fname, fname_prefix, parity_suffix)
        if verbosity.mpb > 0:
            print(f"Outputting {fname}...")

        with h5py.File(fname, "w") as f:
            f["description"] = description.encode()
            self._create_h5_dataset(f, "data")
            self._write_lattice_vectors(f)

            if curfield_type == "n":
                for inv in [False, True]:
                    inv_str = "epsilon_inverse" if inv else "epsilon"
                    for c1 in range(3):
                        for c2 in range(c1, 3):
                            self.mode_solver.get_epsilon_tensor(c1, c2, 0, inv)
                            dataname = f"{inv_str}.{components[c1]}{components[c2]}"
                            self._create_h5_dataset(f, dataname)

                            if with_hermitian_epsilon() and c1 != c2:
                                self.mode_solver.get_epsilon_tensor(c1, c2, 1, inv)
                                dataname += ".i"
                                self._create_h5_dataset(f, dataname)

    def _write_lattice_vectors(self, h5file):
        lattice = np.zeros((3, 3))
        self.mode_solver.get_lattice(lattice)
        h5file["lattice vectors"] = lattice

    def _create_h5_dataset(self, h5file, key):
        h5file[key] = self.get_curfield_as_array(False)

    def _create_fname(self, fname, prefix, parity_suffix):
        parity_str = self.mode_solver.get_parity_string()
        suffix = f".{parity_str}" if parity_suffix and parity_str else ""
        return prefix + fname + suffix + ".h5"

    def compute_field_energy(self):
        return self.mode_solver.compute_field_energy()

    def compute_field_divergence(self):
        return self.mode_solver.compute_field_divergence()

    def compute_energy_in_objects(self, objs):
        return self.mode_solver.compute_energy_in_objects(objs)

    def compute_energy_in_dielectric(self, eps_low, eps_high):
        return self.mode_solver.compute_energy_in_dielectric(eps_low, eps_high)

    def compute_energy_integral(self, f):
        return self.mode_solver.compute_energy_integral(f)

    def compute_field_integral(self, f):
        return self.mode_solver.compute_field_integral(f)

    def compute_group_velocities(self):
        xarg = mp.cartesian_to_reciprocal(mp.Vector3(1), self.geometry_lattice)
        vx = self.mode_solver.compute_group_velocity_component(xarg)
        yarg = mp.cartesian_to_reciprocal(mp.Vector3(y=1), self.geometry_lattice)
        vy = self.mode_solver.compute_group_velocity_component(yarg)
        zarg = mp.cartesian_to_reciprocal(mp.Vector3(z=1), self.geometry_lattice)
        vz = self.mode_solver.compute_group_velocity_component(zarg)

        return [mp.Vector3(x, y, z) for x, y, z in zip(vx, vy, vz)]

    def compute_group_velocity_component(self, direction):
        return self.mode_solver.compute_group_velocity_component(direction)

    def compute_one_group_velocity(self, which_band):
        return self.mode_solver.compute_1_group_velocity(which_band)

    def compute_one_group_velocity_component(self, direction, which_band):
        return self.mode_solver.compute_1_group_velocity_component(
            direction, which_band
        )

    def compute_zparities(self):
        return self.mode_solver.compute_zparities()

    def compute_yparities(self):
        return self.mode_solver.compute_yparities()

    def randomize_fields(self):
        self.mode_solver.randomize_fields()

    def display_kpoint_data(self, name, data):
        if verbosity.mpb > 0:
            k_index = self.mode_solver.get_kpoint_index()
            print(f"{self.parity}{name}:, {k_index}", end="")

            for d in data:
                print(f", {d}", end="")
            print()

    def display_eigensolver_stats(self):
        num_runs = len(self.eigensolver_iters)

        if num_runs <= 0:
            return

        min_iters = min(self.eigensolver_iters)
        max_iters = max(self.eigensolver_iters)
        mean_iters = np.mean(self.eigensolver_iters)
        if verbosity.mpb > 0:
            fmt = "eigensolver iterations for {} kpoints: {}-{}, mean = {}"
            print(fmt.format(num_runs, min_iters, max_iters, mean_iters), end="")

        sorted_iters = sorted(self.eigensolver_iters)
        idx1 = num_runs // 2
        idx2 = ((num_runs + 1) // 2) - 1
        median_iters = 0.5 * (sorted_iters[idx1] + sorted_iters[idx2])
        if verbosity.mpb > 0:
            print(f", median = {median_iters}")

        mean_flops = self.eigensolver_flops / (num_runs * mean_iters)
        if verbosity.mpb > 0:
            print(f"mean flops per iteration = {mean_flops}")

        mean_time = self.total_run_time / (mean_iters * num_runs)
        if verbosity.mpb > 0:
            print(f"mean time per iteration = {mean_time} s")

    def _get_grid_size(self):
        grid_size = mp.Vector3(
            self.resolution[0] * self.geometry_lattice.size.x,
            self.resolution[1] * self.geometry_lattice.size.y,
            self.resolution[2] * self.geometry_lattice.size.z,
        )

        grid_size.x = max(math.ceil(grid_size.x), 1)
        grid_size.y = max(math.ceil(grid_size.y), 1)
        grid_size.z = max(math.ceil(grid_size.z), 1)

        return grid_size

    def _optimize_grid_size(self, grid_size):
        grid_size.x = self.next_factor2357(grid_size.x)
        grid_size.y = self.next_factor2357(grid_size.y)
        grid_size.z = self.next_factor2357(grid_size.z)
        return grid_size

    def next_factor2357(self, n):
        def is_factor2357(n):
            def divby(n, p):
                return divby(n // p, p) if n % p == 0 else n

            return divby(divby(divby(divby(n, 2), 3), 5), 7) == 1

        if is_factor2357(n):
            return n
        return self.next_factor2357(n + 1)

    def init_params(self, p, reset_fields):
        self.mode_solver.init(p, reset_fields, self.geometry, self.default_material)

    def set_parity(self, p):
        self.mode_solver.set_parity(p)

    def solve_kpoint(self, k):
        self.mode_solver.solve_kpoint(k)
        self.freqs = self.get_freqs()

    def run_parity(self, p, reset_fields, *band_functions):
        if self.random_fields and self.randomize_fields not in band_functions:
            band_functions.append(self.randomize_fields)

        start = time.time()

        self.all_freqs = np.zeros((len(self.k_points), self.num_bands))
        self.band_range_data = []

        init_time = time.time()

        if verbosity.mpb > 0:
            print("Initializing eigensolver data")
            print(f"Computing {self.num_bands} bands with {self.tolerance} tolerance")

        self.init_params(p, reset_fields)

        if isinstance(reset_fields, basestring):
            self.load_eigenvectors(reset_fields)

        if verbosity.mpb > 0:
            print(f"{len(self.k_points)} k-points")

            for kp in self.k_points:
                print(f"  {kp}")

            print(f"elapsed time for initialization: {time.time() - init_time}")

        # TODO: Split over multiple processes
        # k_split = list_split(self.k_points, self.k_split_num, self.k_split_index)
        k_split = (0, self.k_points)
        self.mode_solver.set_kpoint_index(k_split[0])

        if self.num_bands > 0:
            for i, k in enumerate(k_split[1]):
                self.current_k = k
                solve_kpoint_time = time.time()
                self.solve_kpoint(k)
                self.iterations = self.mode_solver.get_iterations()
                if verbosity.mpb > 0:
                    print(
                        f"elapsed time for k point: {time.time() - solve_kpoint_time}"
                    )
                self.all_freqs[i, :] = np.array(self.freqs)
                self.band_range_data = self.update_band_range_data(
                    self.band_range_data, self.freqs, k
                )
                self.eigensolver_iters += [self.iterations / self.num_bands]

                for f in band_functions:
                    num_args = get_num_args(f)
                    if num_args == 1:
                        f(self)
                    elif num_args == 2:
                        band = 1
                        while band <= self.num_bands:
                            f(self, band)
                            band += 1
                    else:
                        raise ValueError(
                            "Band function should take 1 or 2 arguments. "
                            "The first must be a ModeSolver instance"
                        )

            if len(k_split[1]) > 1:
                self.output_band_range_data(self.band_range_data)
                self.gap_list = self.output_gaps(self.band_range_data)
            else:
                self.gap_list = []

        end = time.time() - start
        if verbosity.mpb > 0:
            print(f"total elapsed time for run: {end}")
        self.total_run_time += end
        self.eigensolver_flops = self.mode_solver.get_eigensolver_flops()
        self.parity = self.mode_solver.get_parity_string()
        if verbosity.mpb > 0:
            print("done")

    def run(self, *band_functions):
        self.run_parity(mp.NO_PARITY, True, *band_functions)

    def run_zeven(self, *band_functions):
        self.run_parity(mp.EVEN_Z, True, *band_functions)

    def run_zodd(self, *band_functions):
        self.run_parity(mp.ODD_Z, True, *band_functions)

    def run_yeven(self, *band_functions):
        self.run_parity(mp.EVEN_Y, True, *band_functions)

    def run_yodd(self, *band_functions):
        self.run_parity(mp.ODD_Y, True, *band_functions)

    def run_yeven_zeven(self, *band_functions):
        self.run_parity(mp.EVEN_Y + mp.EVEN_Z, True, *band_functions)

    def run_yeven_zodd(self, *band_functions):
        self.run_parity(mp.EVEN_Y + mp.ODD_Z, True, *band_functions)

    def run_yodd_zeven(self, *band_functions):
        self.run_parity(mp.ODD_Y + mp.EVEN_Z, True, *band_functions)

    def run_yodd_zodd(self, *band_functions):
        self.run_parity(mp.ODD_Y + mp.ODD_Z, True, *band_functions)

    run_te = run_zeven
    run_tm = run_zodd
    run_te_yeven = run_yeven_zeven
    run_te_yodd = run_yodd_zeven
    run_tm_yeven = run_yeven_zodd
    run_tm_yodd = run_yodd_zodd

    def find_k(
        self,
        p,
        omega,
        band_min,
        band_max,
        korig_and_kdir,
        tol,
        kmag_guess,
        kmag_min,
        kmag_max,
        *band_funcs,
    ):

        num_bands_save = self.num_bands
        kpoints_save = self.k_points
        nb = band_max - band_min + 1
        kdir = korig_and_kdir[1] if type(korig_and_kdir) is list else korig_and_kdir
        lat = self.geometry_lattice
        kdir1 = mp.cartesian_to_reciprocal(
            mp.reciprocal_to_cartesian(kdir, lat).unit(), lat
        )

        if type(korig_and_kdir) is list:
            korig = korig_and_kdir[0]
        else:
            korig = mp.Vector3()

        # k0s is a list caching the best k value found for each band:
        if type(kmag_guess) is list:
            k0s = kmag_guess
        else:
            k0s = [kmag_guess] * (band_max - band_min + 1)

        # dict to memoize all "band: k" results
        bktab = {}

        def rootfun(b):
            def _rootfun(k):
                # First, look in the cached table
                tab_val = bktab.get((b, k), None)
                if tab_val:
                    if verbosity.mpb > 0:
                        print(f"find-k {b} at {k}: {tab_val[0]} (cached)")
                    return tab_val
                else:
                    self.num_bands = b
                    self.k_points = [korig + kdir1.scale(k)]
                    self.run_parity(p, False)
                    v = self.mode_solver.compute_group_velocity_component(kdir1)

                    # Cache computed values
                    for _b, _f, _v in zip(
                        range(band_min, b - band_min + 1),
                        self.freqs[band_min - 1 :],
                        v[band_min - 1 :],
                    ):
                        tabval = bktab.get((_b, k0s[_b - band_min]), None)

                        if not tabval or abs(_f - omega) < abs(tabval[0]):
                            k0s[_b - band_min + 1] = k

                        bktab[(_b, k)] = (_f - omega, _v)

                    fun = self.freqs[-1] - omega
                    if verbosity.mpb > 0:
                        print(f"find-k {b} at {k}: {fun}")
                    return (fun, v[-1])

            return _rootfun

        # Don't let previous computations interfere
        if self.mode_solver:
            self.randomize_fields()

        ks = []
        for b in range(band_max, band_max - nb, -1):
            ks.append(
                mp.find_root_deriv(
                    rootfun(b), tol, kmag_min, kmag_max, k0s[b - band_min]
                )
            )

        if band_funcs:
            for b, k in zip(range(1, band_max + 1), reversed(ks)):
                self.num_bands = b
                self.k_points = [korig + kdir1.scale(k)]

                def bfunc(ms, b_prime):
                    if b_prime == b:
                        for f in band_funcs:
                            apply_band_func_thunk(ms, f, b, True)

                self.run_parity(p, False, bfunc)

        self.num_bands = num_bands_save
        self.k_points = kpoints_save
        ks = list(reversed(ks))
        if verbosity.mpb > 0:
            print(
                f"{self.parity}kvals:, {omega}, {band_min}, {band_max}",
                end="",
            )
            for k in korig:
                print(f", {k}", end="")
            for k in kdir1:
                print(f", {k}", end="")
            for k in ks:
                print(f", {k}", end="")
            print()

        return ks

    def first_brillouin_zone(self, k):
        """
        Function to convert a k-point k into an equivalent point in the
        first Brillouin zone (not necessarily the irreducible Brillouin zone)
        """

        def n(k):
            return mp.reciprocal_to_cartesian(k, self.geometry_lattice).norm()

        def try_plus(k, v):
            return try_plus(k + v, v) if n(k + v) < n(k) else k

        def _try(k, v):
            return try_plus(try_plus(k, v), mp.Vector3() - v)

        try_list = [
            mp.Vector3(1, 0, 0),
            mp.Vector3(0, 1, 0),
            mp.Vector3(0, 0, 1),
            mp.Vector3(0, 1, 1),
            mp.Vector3(1, 0, 1),
            mp.Vector3(1, 1, 0),
            mp.Vector3(0, 1, -1),
            mp.Vector3(1, 0, -1),
            mp.Vector3(1, -1, 0),
            mp.Vector3(1, 1, 1),
            mp.Vector3(-1, 1, 1),
            mp.Vector3(1, -1, 1),
            mp.Vector3(1, 1, -1),
        ]

        def try_all(k):
            return functools.reduce(_try, try_list, k)

        def try_all_and_repeat(k):
            knew = try_all(k)
            return try_all_and_repeat(knew) if n(knew) < n(k) else k

        k0 = k - mp.Vector3(*[round(x) for x in k])

        return try_all_and_repeat(k0) if n(k0) < n(k) else try_all_and_repeat(k)

    def get_dominant_planewave(self, band):
        return self.mode_solver.get_dominant_planewave(band)

    def transformed_overlap(self, W, w):
        return self.mode_solver.transformed_overlap(W, w)

    def compute_symmetry(self, band, W, w):
        return self.mode_solver.compute_symmetry(band, W, w)

    def compute_symmetries(self, W, w):
        return [
            self.mode_solver.compute_symmetry(band, W, w)
            for band in range(1, self.num_bands + 1)
        ]


# Predefined output functions (functions of the band index), for passing to `run`


def output_hfield(ms, which_band):
    ms.get_hfield(which_band, False)
    ms.output_field()


def output_hfield_x(ms, which_band):
    ms.get_hfield(which_band, False)
    ms.output_field_x()


def output_hfield_y(ms, which_band):
    ms.get_hfield(which_band, False)
    ms.output_field_y()


def output_hfield_z(ms, which_band):
    ms.get_hfield(which_band, False)
    ms.output_field_z()


def output_bfield(ms, which_band):
    ms.get_bfield(which_band, False)
    ms.output_field()


def output_bfield_x(ms, which_band):
    ms.get_bfield(which_band, False)
    ms.output_field_x()


def output_bfield_y(ms, which_band):
    ms.get_bfield(which_band, False)
    ms.output_field_y()


def output_bfield_z(ms, which_band):
    ms.get_bfield(which_band, False)
    ms.output_field_z()


def output_dfield(ms, which_band):
    ms.get_dfield(which_band, False)
    ms.output_field()


def output_dfield_x(ms, which_band):
    ms.get_dfield(which_band, False)
    ms.output_field_x()


def output_dfield_y(ms, which_band):
    ms.get_dfield(which_band, False)
    ms.output_field_y()


def output_dfield_z(ms, which_band):
    ms.get_dfield(which_band, False)
    ms.output_field_z()


def output_efield(ms, which_band):
    ms.get_efield(which_band, False)
    ms.output_field()


def output_efield_x(ms, which_band):
    ms.get_efield(which_band, False)
    ms.output_field_x()


def output_efield_y(ms, which_band):
    ms.get_efield(which_band, False)
    ms.output_field_y()


def output_efield_z(ms, which_band):
    ms.get_efield(which_band, False)
    ms.output_field_z()


def output_bpwr(ms, which_band):
    ms.get_bfield(which_band, False)
    ms.compute_field_energy()
    ms.output_field()


def output_dpwr(ms, which_band):
    ms.get_dfield(which_band, False)
    ms.compute_field_energy()
    ms.output_field()


def output_tot_pwr(ms, which_band):
    ms.get_tot_pwr(which_band)
    ms.output_field_to_file(-1, f"{ms.get_filename_prefix()}tot.")


def output_dpwr_in_objects(output_func, min_energy, objects=[]):
    """
    The following function returns an output function that calls output_func for
    bands with D energy in objects > min-energy. For example,
    output_dpwr_in_objects(output_dfield, 0.20, some_object) would return an
    output function that would spit out the D field for bands with at least %20
    of their D energy in some-object.
    """

    def _output(ms, which_band):
        ms.get_dfield(which_band, False)
        ms.compute_field_energy()
        energy = ms.compute_energy_in_objects(objects)
        if verbosity.mpb > 0:
            fmt = "dpwr:, {}, {}, {} "
            print(fmt.format(which_band, ms.freqs[which_band - 1], energy))
        if energy >= min_energy:
            apply_band_func(ms, output_func, which_band)

    return _output


def output_charge_density(ms, which_band):
    ms.get_charge_density(which_band)
    ms.output_field_to_file(-1, ms.get_filename_prefix())


def output_poynting(ms, which_band):
    ms.get_poynting(which_band)
    ms.output_field_to_file(-1, f"{ms.get_filename_prefix()}flux.")


def output_poynting_x(ms, which_band):
    ms.get_poynting(which_band)
    ms.output_field_to_file(0, f"{ms.get_filename_prefix()}flux.")


def output_poynting_y(ms, which_band):
    ms.get_poynting(which_band)
    ms.output_field_to_file(1, f"{ms.get_filename_prefix()}flux.")


def output_poynting_z(ms, which_band):
    ms.get_poynting(which_band)
    ms.output_field_to_file(2, f"{ms.get_filename_prefix()}flux.")


def display_yparities(ms):
    ms.display_kpoint_data("yparity", ms.mode_solver.compute_yparities())


def display_zparities(ms):
    ms.display_kpoint_data("zparity", ms.mode_solver.compute_zparities())


def display_group_velocities(ms):
    ms.display_kpoint_data("velocity", ms.compute_group_velocities())


# Band functions to pick a canonical phase for the eigenstate of the
# given band based upon the spatial representation of the given field


def fix_hfield_phase(ms, which_band):
    ms.get_hfield(which_band, False)
    ms.mode_solver.fix_field_phase()


def fix_bfield_phase(ms, which_band):
    ms.get_bfield(which_band, False)
    ms.mode_solver.fix_field_phase()


def fix_dfield_phase(ms, which_band):
    ms.get_dfield(which_band, False)
    ms.mode_solver.fix_field_phase()


def fix_efield_phase(ms, which_band):
    ms.get_efield(which_band, False)
    ms.mode_solver.fix_field_phase()


def apply_band_func_thunk(ms, band_func, which_band, eval_thunk):
    """
    We need a special function to evaluate band functions, since band functions
    can either be a function of the band number or a thunk (function of no arguments,
    evaluated once per k-point).
    """
    if get_num_args(band_func) == 1:
        if eval_thunk:
            band_func(ms)  # evaluate thunks once per k-point
    else:
        band_func(ms, which_band)


def apply_band_func(ms, band_func, which_band):
    apply_band_func_thunk(ms, band_func, which_band, which_band == 1)


def combine_band_functions(*band_funcs):
    """Combines zero or more band functions into one"""

    def _combine(ms, which_band):
        for f in band_funcs:
            apply_band_func(ms, f, which_band)

    return _combine


def output_at_kpoint(kpoint, *band_funcs):
    """Only invoke the given band functions for the specified k-point"""

    band_func = combine_band_functions(*band_funcs)

    def _output_at_kpoint(ms, which_band):
        if ms.current_k.close(kpoint, tol=1e-8 * kpoint.norm()):
            band_func(ms, which_band)

    return _output_at_kpoint