Compute the defect densities using dilute solution model.

    Args:
        structure: pymatgen.core.structure.Structure object representing the
            primitive or unitcell of the crystal.
        e0: The total energy of the undefected system.
            This is E0 from VASP calculation.
        vac_defs: List of vacancy defect parameters in the dictionary format.
            The keys of the dict associated with each vacancy defect are
            1) site_index, 2) site_specie, 3) site_multiplicity, and
            4) energy. 1-3 can be obtained from
            pymatgen.analysis.defects.point_defects.Vacancy class.
            Site index is expected to start with 1 (fortran index).
        antisite_defs: List of antisite defect parameters in the dictionary
            format. The keys of the dict associated with each antisite defect
            are 1) site_index, 2) site_specie, 3) site_multiplicity,
            4) substitution_specie, and 5) energy. 1-3 can be obtained
            from pymatgen.analysis.defects.point_defects.Vacancy class.
        T: Temperature in Kelvin
        trial_chem_pot (optional): Trial chemical potentials to speedup
            the plot generation. Format is {el1:mu1,...}
        generate (string): Options are plot or energy
            Chemical potentials are also returned with energy option.
            If energy option is not chosen, plot is generated.

    Returns:
        If generate=plot, the plot data is generated and returned in
        HighCharts format.
        If generate=energy, defect formation enthalpies and chemical
        potentials are returned.

def dilute_solution_model(structure, e0, vac_defs, antisite_defs, T, trial_chem_pot=None, generate='plot'):
    """
    Compute the defect densities using dilute solution model.

    Args:
        structure: pymatgen.core.structure.Structure object representing the
            primitive or unitcell of the crystal.
        e0: The total energy of the undefected system.
            This is E0 from VASP calculation.
        vac_defs: List of vacancy defect parameters in the dictionary format.
            The keys of the dict associated with each vacancy defect are
            1) site_index, 2) site_specie, 3) site_multiplicity, and
            4) energy. 1-3 can be obtained from
            pymatgen.analysis.defects.point_defects.Vacancy class.
            Site index is expected to start with 1 (fortran index).
        antisite_defs: List of antisite defect parameters in the dictionary
            format. The keys of the dict associated with each antisite defect
            are 1) site_index, 2) site_specie, 3) site_multiplicity,
            4) substitution_specie, and 5) energy. 1-3 can be obtained
            from pymatgen.analysis.defects.point_defects.Vacancy class.
        T: Temperature in Kelvin
        trial_chem_pot (optional): Trial chemical potentials to speedup
            the plot generation. Format is {el1:mu1,...}
        generate (string): Options are plot or energy
            Chemical potentials are also returned with energy option.
            If energy option is not chosen, plot is generated.

    Returns:
        If generate=plot, the plot data is generated and returned in
        HighCharts format.
        If generate=energy, defect formation enthalpies and chemical
        potentials are returned.
    """

    if not check_input(vac_defs):
        raise ValueError('Vacancy energy is not defined')
    if not check_input(antisite_defs):
        raise ValueError('Antisite energy is not defined')

    formation_energies = {}
    formation_energies['vacancies'] = copy.deepcopy(vac_defs)
    formation_energies['antisites'] = copy.deepcopy(antisite_defs)
    for vac in formation_energies['vacancies']:
        del vac['energy']
    for asite in formation_energies['antisites']:
        del asite['energy']
    # Setup the system
    site_species = [vac_def['site_specie'] for vac_def in vac_defs]
    multiplicity = [vac_def['site_multiplicity'] for vac_def in vac_defs]
    m = len(set(site_species))  # distinct species
    n = len(vac_defs)  # inequivalent sites

    # Reduce the system and associated parameters such that only distinctive
    # atoms are retained
    comm_div = gcd(*tuple(multiplicity))
    multiplicity = [val / comm_div for val in multiplicity]
    e0 = e0 / comm_div
    T = Float(T)

    #c0 = np.diag(multiplicity)
    c0 = np.diag(np.ones(n))
    mu = [Symbol('mu' + i.__str__()) for i in range(m)]

    # Generate maps for hashing
    # Generate specie->mu map and use it for site->mu map
    specie_order = []  # Contains hash for site->mu map    Eg: [Al, Ni]
    site_specie_set = set()  # Eg: {Ni, Al}
    for i in range(n):
        site_specie = site_species[i]
        if site_specie not in site_specie_set:
            site_specie_set.add(site_specie)
            specie_order.append(site_specie)
    site_mu_map = []  # Eg: [mu0,mu0,mu0,mu1] where mu0->Al, and mu1->Ni
    for i in range(n):
        site_specie = site_species[i]
        j = specie_order.index(site_specie)
        site_mu_map.append(j)
    specie_site_index_map = []  # Eg: [(0,3),(3,4)] for Al & Ni
    for i in range(m):
        low_ind = site_species.index(specie_order[i])
        if i < m - 1:
            hgh_ind = site_species.index(specie_order[i + 1])
        else:
            hgh_ind = n
        specie_site_index_map.append((low_ind, hgh_ind))
    """
    dC: delta concentration matrix:
    dC[i,j,k]: Concentration change of atom i, due to presence of atom
    j on lattice site k
    Special case is [i,i,i] which is considered as vacancy
    Few cases: dC[i,i,i] = -1 due to being vacancy special case
                dC[k,k,i] = +1 due to increment in k at i lattice if i
                               lattice type is of different element
                dC[i,k,i] = -1 due to decrement of ith type atom due to
                presence of kth type atom on ith sublattice and kth type
                atom specie is different from ith sublattice atom specie
                dC[i,k,k] = 0 due to no effect on ith type atom
                dC[i,j,k] = 0 if i!=j!=k
    """
    dC = np.zeros((n, n, n), dtype=np.int)
    for i in range(n):
        for j in range(n):
            for k in range(n):
                if i == j and site_species[j] != site_species[k] and \
                        site_species[i] != site_species[k]:
                    dC[i, j, k] = 1
        for j in range(n):
            for k in range(n):
                if i == k:
                    dC[i, j, k] = -1
    for k in range(n):
        for j in range(n):
            for i in range(n):
                if i != j:
                    if site_species[j] == site_species[k]:
                        dC[i, j, k] = 0

    for ind_map in specie_site_index_map:
        if ind_map[1] - ind_map[0] > 1:
            for index1 in range(ind_map[0] + 1, ind_map[1]):
                for index2 in range(ind_map[0]):
                    for i in range(n):
                        dC[i, index1, index2] = 0
                for index2 in range(ind_map[1], n):
                    for i in range(n):
                        dC[i, index1, index2] = 0

    # dE matrix: Flip energies (or raw defect energies)
    els = [vac_def['site_specie'] for vac_def in vac_defs]
    dE = []
    for i in range(n):
        dE.append([])
    for i in range(n):
        for j in range(n):
            dE[i].append(0)

    for j in range(n):
        for i in range(n):
            if i == j:
                dE[i][j] = vac_defs[i]['energy']
            else:
                sub_specie = vac_defs[i]['site_specie']
                site_specie = vac_defs[j]['site_specie']
                if site_specie == sub_specie:
                    dE[i][j] = 0
                else:
                    for as_def in antisite_defs:
                        if int(as_def['site_index']) == j + 1 and \
                                sub_specie == as_def['substitution_specie']:
                            dE[i][j] = as_def['energy']
                            break
    dE = np.array(dE)

    # Initialization for concentrations
    # c(i,p) == presence of ith type atom on pth type site
    c = Matrix(n, n, [0] * n**2)
    for i in range(n):
        for p in range(n):
            c[i, p] = Integer(c0[i, p])
            site_flip_contribs = []
            for epi in range(n):
                sum_mu = sum([mu[site_mu_map[j]] * Integer(dC[j, epi, p]) for j in range(n)])
                flip = Integer(dC[i, epi, p]) * \
                    exp(-(dE[epi, p] - sum_mu) / (k_B * T))
                if flip not in site_flip_contribs:
                    site_flip_contribs.append(flip)
                    c[i, p] += flip

    total_c = []
    for ind in specie_site_index_map:
        val = 0
        for i in range(*ind):
            sum_i = sum([c[i, j] * multiplicity[j] for j in range(n)])
            val += sum_i
        total_c.append(val)

    c_ratio = [total_c[-1] / total_c[i] for i in range(m)]

    # Expression for Omega, the Grand Potential
    omega1 = e0 - sum([mu[site_mu_map[i]] * sum(c0[i, :]) * multiplicity[i] for i in range(n)])
    omega2 = []
    fm_en_eff = []
    used_dEs = []
    for p_r in range(n):
        for epi in range(n):
            sum_mu = sum([mu[site_mu_map[j]] * dC[j, epi, p_r] for j in range(n)])
            if p_r != epi and site_mu_map[p_r] == site_mu_map[epi]:
                continue
            if dE[epi, p_r] not in used_dEs:
                omega2.append(k_B * T * multiplicity[p_r] * exp(-(dE[epi, p_r] - sum_mu) / (k_B * T)))
                fm_en_eff.append(dE[epi, p_r] - sum_mu)
                used_dEs.append(dE[epi, p_r])
    omega = omega1 - sum(omega2)

    # Compute composition range
    li = specie_site_index_map[0][0]
    hi = specie_site_index_map[0][1]
    comp1_min = sum(multiplicity[li:hi]) / sum(multiplicity) * 100 - 1
    comp1_max = sum(multiplicity[li:hi]) / sum(multiplicity) * 100 + 1
    delta = float(comp1_max - comp1_min) / 120.0
    yvals = []
    for comp1 in np.arange(comp1_min, comp1_max + delta, delta):
        comp2 = 100 - comp1
        y = comp2 / comp1
        yvals.append(y)

    def reduce_mu():
        omega = [e0 - sum([mu[site_mu_map[i]] * sum(c0[i, :]) for i in range(n)])]
        x = solve(omega)
        return x

    def compute_mus_by_search():
        # Compute trial mu
        mu_red = reduce_mu()

        mult = multiplicity
        specie_concen = [sum(mult[ind[0]:ind[1]]) for ind in specie_site_index_map]
        y_vect = [specie_concen[-1] / specie_concen[i] for i in range(m)]
        vector_func = [y_vect[i] - c_ratio[i] for i in range(m - 1)]
        vector_func.append(omega)
        min_diff = 1e10
        mu_vals = None
        c_val = None
        m1_min = -20.0
        if e0 > 0:
            m1_max = 10  # Search space needs to be modified
        else:
            m1_max = 0
        for m1 in np.arange(m1_min, m1_max, 0.01):
            m0 = mu_red[mu[0]].subs(mu[-1], m1)

            try:
                x = nsolve(vector_func, mu, [m0, m1], module="numpy")
            except:
                continue

            c_val = c.subs(dict(zip(mu, x)))
            # if all(x >= 0 for x in c_val):
            specie_concen = []
            for ind in specie_site_index_map:
                specie_concen.append(sum([sum(c_val[i, :]) for i in range(*ind)]))
            y_comp = [specie_concen[-1] / specie_concen[i] for i in range(m)]
            diff = math.sqrt(sum([pow(abs(y_comp[i] - y_vect[i]), 2) for i in range(m)]))
            if diff < min_diff:
                min_diff = diff
                mu_vals = x

        if mu_vals:
            mu_vals = [float(mu_val) for mu_val in mu_vals]
        else:
            raise ValueError()
        return mu_vals

    def compute_def_formation_energies():
        i = 0
        for vac_def in vac_defs:
            site_specie = vac_def['site_specie']
            ind = specie_order.index(site_specie)
            uncor_energy = vac_def['energy']
            formation_energy = uncor_energy + mu_vals[ind]
            formation_energies['vacancies'][i]['formation_energy'] = formation_energy
            specie_ind = site_mu_map[i]
            indices = specie_site_index_map[specie_ind]
            specie_ind_del = indices[1] - indices[0]
            cur_ind = i - indices[0] + 1
            if not specie_ind_del - 1:
                label = '$V_{' + site_specie + '}$'
            else:
                label = '$V_{' + site_specie + '_' + str(cur_ind) + '}$'
            formation_energies['vacancies'][i]['label'] = label
            i += 1
        i = 0
        for as_def in antisite_defs:
            site_specie = as_def['site_specie']
            sub_specie = as_def['substitution_specie']
            ind1 = specie_order.index(site_specie)
            ind2 = specie_order.index(sub_specie)
            uncor_energy = as_def['energy']
            formation_energy = uncor_energy + mu_vals[ind1] - mu_vals[ind2]
            formation_energies['antisites'][i]['formation_energy'] = formation_energy
            specie_ind = site_mu_map[i]
            indices = specie_site_index_map[specie_ind]
            specie_ind_del = indices[1] - indices[0]
            cur_ind = i - indices[0] + 1
            if not specie_ind_del - 1:
                label = '$' + sub_specie + '_{' + site_specie + '}$'
            else:
                label = '$' + sub_specie + '_{' + site_specie + '_' + str(cur_ind) + '}$'
            formation_energies['antisites'][i]['label'] = label
            i += 1
        return formation_energies

    # If generate option is energy compute effective formation energies
    # at ideal stoichiometry and return the formation energies and chem pot.
    if generate == 'energy':
        if not trial_chem_pot:
            mu_vals = compute_mus_by_search()
        else:
            try:
                mu_vals = [trial_chem_pot[element] for element in specie_order]
            except:
                mu_vals = compute_mus()

        formation_energies = compute_def_formation_energies()
        mu_dict = dict(zip(specie_order, mu_vals))
        return formation_energies, mu_dict

    if not trial_chem_pot:
        # Try computing mus by assuming one of the defects is dominant at 0.01
        # concen.  First vacancy is tried and then antisite

        # Generate trial mus assuming vacancy as dominant defect
        # for specie-0 at lower yval
        li = specie_site_index_map[0][0]
        hi = specie_site_index_map[0][1]
        li1 = specie_site_index_map[1][0]
        hi1 = specie_site_index_map[1][1]
        spec_mult = [sum(multiplicity[li:hi]), sum(multiplicity[li1:hi1])]
        ln_def_conc = 4.60517
        for i in range(li, hi):
            vac_flip_en = vac_defs[i]['energy']
            mu_vals = [ln_def_conc * k_B * T - vac_flip_en]
            mu_vals.append((e0 - spec_mult[0] * mu_vals[0]) / spec_mult[1])
            comp_ratio = yvals[0]

            # Test if the trial mus are good
            vector_func = [comp_ratio - c_ratio[0]]
            vector_func.append(omega)
            try:
                mu_vals = nsolve(vector_func, mu, mu_vals)
                if mu_vals:
                    mu_vals = [float(mu_val) for mu_val in mu_vals]
                break
            except:  # Go for antisite as dominant defect
                mu_gs = [Symbol('mu_gs' + j.__str__()) for j in range(m)]

                eqs = [mu_gs[0] - mu_gs[1] - (ln_def_conc * k_B * T - antisite_defs[i]['energy'])]
                eqs.append(spec_mult[0] * mu_gs[0] + spec_mult[1] * mu_gs[1] - e0)
                x = solve(eqs, mu_gs)
                #mu_names = sorted([key.name for key in x.keys()])
                mu_vals = []
                for key in sorted(x.keys(), key=lambda inp: inp.name):
                    mu_vals.append(x[key])
                vector_func = [comp_ratio - c_ratio[0]]
                vector_func.append(omega)

                try:
                    mu_vals = nsolve(vector_func, mu, mu_vals)
                    if mu_vals:
                        mu_vals = [float(mu_val) for mu_val in mu_vals]
                    break
                except:  # Go to the default option (search the space)
                    pass
        else:
            mu_vals = compute_mus_by_search()

    else:
        try:
            mu_vals = [trial_chem_pot[element] for element in specie_order]
        except:
            mu_vals = compute_mus_by_search()

    # Compile mu's for all composition ratios in the range
    #+/- 1% from the stoichiometry
    result = {}
    i = 0
    len_y = len(yvals)
    failed_y, failed_i = [], []
    for y in yvals:
        vector_func = [y - c_ratio[0]]
        vector_func.append(omega)
        try:
            x = nsolve(vector_func, mu, mu_vals, module="numpy")
            if x:
                mu_vals = [float(mu_val) for mu_val in x]
        except:
            failed_y.append(y)
            failed_i.append(i)
            continue
        result[y] = list(mu_vals)
        x = None
        i += 1

    def get_next_mu_val(i):
        if i >= len(yvals):
            return None

        y = yvals[i + 1]
        x = result.get(y, None)
        if x:
            mu_vals = [float(mu_val) for mu_val in x]
            return mu_vals
        else:
            return get_next_mu_val(i + 1)

    def get_prev_mu_val(i):
        if i <= 0:
            return None

        y = yvals[i - 1]
        x = result.get(y, None)
        if x:
            mu_vals = [float(mu_val) for mu_val in x]
            return mu_vals
        else:
            return get_next_mu_val(i - 1)

    # Try to get better trial mus for failed cases
    for j in range(len(failed_y)):
        i = failed_i[j]

        prev_mu_val = get_prev_mu_val(i)
        if not prev_mu_val:
            continue
        next_mu_val = get_next_mu_val(i)
        if not next_mu_val:
            continue

        y = failed_y[j]
        vector_func = [y - c_ratio[0]]
        vector_func.append(omega)
        trial_mu = list(map(lambda x: float(sum(x)) / len(x), zip(prev_mu_val, next_mu_val)))
        try:
            x = nsolve(vector_func, mu, trial_mu, module="numpy")
            if x:
                mu_vals = [float(mu_val) for mu_val in x]
        except:
            continue
        result[y] = mu_vals
        x = None

    # Alternate way of calculating trial mus for failed cases
    # by taking average of trial mus at extremes.
    # for j in range(len(failed_y)):
    #    y = yvals[0]
    #    prev_mu_val = result[y]
    #    y = yvals[-1]
    #    next_mu_val = result[y]
    #
    #    trial_mu = list(map(lambda x: float(sum(x))/len(x),  \
    #            zip(prev_mu_val,next_mu_val)))
    #    y = failed_y[j]
    #    vector_func = [y-c_ratio[0]]
    #    vector_func.append(omega)
    #    try:
    #        x = nsolve(vector_func,mu,trial_mu,module="numpy")
    #        if x:
    #            mu_vals = [float(mu_val) for mu_val in x]
    #    except:
    #        continue
    #    result[y] = list(mu_vals)

    if len(result.keys()) < len(yvals) / 2:
        raise ValueError('Not sufficient data')

    res = []
    new_mu_dict = {}
    # Compute the concentrations for all the compositions
    for key in sorted(result.keys()):
        mu_val = result[key]
        total_c_val = [total_c[i].subs(dict(zip(mu, mu_val))) for i in range(len(total_c))]
        c_val = c.subs(dict(zip(mu, mu_val)))
        res1 = []
        # Concentration of first element/over total concen
        res1.append(float(total_c_val[0] / sum(total_c_val)))
        new_mu_dict[res1[0]] = mu_val
        sum_c0 = sum([c0[i, i] for i in range(n)])
        for i in range(n):
            for j in range(n):
                if i == j:  # Vacancy
                    vac_conc = float(exp(-(mu_val[site_mu_map[i]] + dE[i, i]) / (k_B * T)))
                    res1.append(vac_conc)
                else:  # Antisite
                    res1.append(float(c_val[i, j] / c0[j, j]))
        res.append(res1)

    res = np.array(res)
    dtype = [(str('x'), np.float64)] + [(str('y%d%d' % (i, j)), np.float64) for i in range(n) for j in range(n)]
    res1 = np.sort(res.view(dtype), order=[str('x')], axis=0)

    conc_data = {}
    """Because all the plots have identical x-points storing it in a
    single array"""
    conc_data['x'] = [dat[0][0] for dat in res1]  # x-axis data
    # Element whose composition is varied. For x-label
    conc_data['x_label'] = els[0] + " mole fraction"
    conc_data['y_label'] = "Point defect concentration"
    conc = []
    for i in range(n):
        conc.append([])
        for j in range(n):
            conc[i].append([])
    for i in range(n):
        for j in range(n):
            y1 = [dat[0][i * n + j + 1] for dat in res1]
            conc[i][j] = y1

    y_data = []
    for i in range(n):
        data = conc[i][i]
        specie = els[i]
        specie_ind = site_mu_map[i]
        indices = specie_site_index_map[specie_ind]
        specie_ind_del = indices[1] - indices[0]
        cur_ind = i - indices[0] + 1
        vac_string = "$Vac_{"
        if not specie_ind_del - 1:
            label = vac_string + specie + '}$'
        else:
            label = vac_string + specie + '_' + str(cur_ind) + '}$'
        # Plot data and legend info
        y_data.append({'data': data, 'name': label})

    for i in range(n):
        site_specie = els[i]
        specie_ind = site_mu_map[i]
        indices = specie_site_index_map[specie_ind]
        specie_ind_del = indices[1] - indices[0]
        cur_ind = i - indices[0] + 1
        for j in range(m):  # Antisite plot dat
            sub_specie = specie_order[j]
            if sub_specie == site_specie:
                continue
            if not specie_ind_del - 1:
                label = '$' + sub_specie + '_{' + site_specie + '}$'
            else:
                label = '$' + sub_specie + '_{' + site_specie + '_' + str(cur_ind) + '}$'
            inds = specie_site_index_map[j]
            # TODO: Investigate the value below
            data = np.sum([conc[ind][i] for ind in range(*inds)], axis=0)
            data = data.tolist()
            y_data.append({'data': data, 'name': label})

    conc_data['y'] = y_data

    # Compute the  formation energies
    def compute_vac_formation_energies(mu_vals):
        en = []
        for vac_def in vac_defs:
            site_specie = vac_def['site_specie']
            ind = specie_order.index(site_specie)
            uncor_energy = vac_def['energy']
            formation_energy = uncor_energy + mu_vals[ind]
            en.append(float(formation_energy))
        return en

    en_res = []
    for key in sorted(new_mu_dict.keys()):
        mu_val = new_mu_dict[key]
        en_res.append(compute_vac_formation_energies(mu_val))

    en_data = {'x_label': els[0] + ' mole fraction', 'x': []}
    en_data['x'] = [dat[0][0] for dat in res1]  # x-axis data

    i = 0
    y_data = []
    for vac_def in vac_defs:
        data = [data[i] for data in en_res]
        site_specie = vac_def['site_specie']
        ind = specie_order.index(site_specie)
        specie_ind = site_mu_map[i]
        indices = specie_site_index_map[specie_ind]
        specie_ind_del = indices[1] - indices[0]
        cur_ind = i - indices[0] + 1
        vac_string = "$Vac_{"
        if not specie_ind_del - 1:
            label = vac_string + site_specie + '}$'
        else:
            label = vac_string + site_specie + '_' + str(cur_ind) + '}$'
        y_data.append({'data': data, 'name': label})
        i += 1

    def compute_as_formation_energies(mu_vals):
        en = []
        for as_def in antisite_defs:
            site_specie = as_def['site_specie']
            sub_specie = as_def['substitution_specie']
            ind1 = specie_order.index(site_specie)
            ind2 = specie_order.index(sub_specie)
            uncor_energy = as_def['energy']
            form_en = uncor_energy + mu_vals[ind1] - mu_vals[ind2]
            en.append(form_en)
        return en

    en_res = []
    for key in sorted(new_mu_dict.keys()):
        mu_val = new_mu_dict[key]
        en_res.append(compute_as_formation_energies(mu_val))
    i = 0
    for as_def in antisite_defs:
        data = [data[i] for data in en_res]
        site_specie = as_def['site_specie']
        sub_specie = as_def['substitution_specie']
        ind1 = specie_order.index(site_specie)
        ind2 = specie_order.index(sub_specie)
        specie_ind = site_mu_map[i]
        indices = specie_site_index_map[specie_ind]
        specie_ind_del = indices[1] - indices[0]
        cur_ind = i - indices[0] + 1
        if not specie_ind_del - 1:
            label = '$' + sub_specie + '_{' + site_specie + '}$'
        else:
            label = '$' + sub_specie + '_{' + site_specie + '_' + str(cur_ind) + '}$'
        y_data.append({'data': data, 'name': label})
        i += 1

    en_data['y'] = y_data

    # Return chem potential as well
    mu_data = {'x_label': els[0] + ' mole fraction', 'x': []}
    mu_data['x'] = [dat[0][0] for dat in res1]  # x-axis data

    y_data = []
    for j in range(m):
        specie = specie_order[j]
        mus = [new_mu_dict[key][j] for key in sorted(new_mu_dict.keys())]
        y_data.append({'data': mus, 'name': specie})
    mu_data['y'] = y_data

    return conc_data, en_data, mu_data