#include 
#include 
#include 
#include 
#include "source_base/formatter.h"
#include "source_io/cif_io.h"
#include 
#include 
#include "source_base/tool_quit.h"
#ifdef __MPI
#include "source_base/parallel_common.h"
#endif

double deg2rad(double deg) { return deg * M_PI / 180.0; }
double rad2deg(double rad) { return rad * 180.0 / M_PI; }

void _build_chem_formula(const int natom, 
                         const std::string* atom_site_labels,
                         std::string& sum,
                         std::string& structural)
{
    sum.clear();
    structural.clear();
    std::vector kinds;
    std::vector labels(natom);
    std::copy(atom_site_labels, atom_site_labels + natom, labels.begin());
    for (int i = 0; i < natom; ++i)
    {
        if (std::find(kinds.begin(), kinds.end(), labels[i]) == kinds.end())
        {
            kinds.push_back(labels[i]);
        }
    }
    std::vector counts(kinds.size());
    std::transform(kinds.begin(), kinds.end(), counts.begin(), [&labels](const std::string& kind) {
        return std::count(labels.begin(), labels.end(), kind);
    });
    for (size_t i = 0; i < kinds.size(); ++i)
    {
        sum += kinds[i];
        structural += kinds[i];
        if (counts[i] > 1)
        {
            sum += std::to_string(counts[i]);
        }
    }
}

void vec_to_abc_angles(const double* vec, double* abc_angles)
{
    const std::vector a = {vec[0], vec[1], vec[2]};
    const std::vector b = {vec[3], vec[4], vec[5]};
    const std::vector c = {vec[6], vec[7], vec[8]};
    const double anorm = std::sqrt(a[0] * a[0] + a[1] * a[1] + a[2] * a[2]);
    const double bnorm = std::sqrt(b[0] * b[0] + b[1] * b[1] + b[2] * b[2]);
    const double cnorm = std::sqrt(c[0] * c[0] + c[1] * c[1] + c[2] * c[2]);
    const double alpha = std::acos((b[0] * c[0] + b[1] * c[1] + b[2] * c[2]) / (bnorm * cnorm));
    const double beta = std::acos((a[0] * c[0] + a[1] * c[1] + a[2] * c[2]) / (anorm * cnorm));
    const double gamma = std::acos((a[0] * b[0] + a[1] * b[1] + a[2] * b[2]) / (anorm * bnorm));
    abc_angles[0] = anorm;
    abc_angles[1] = bnorm;
    abc_angles[2] = cnorm;
    abc_angles[3] = rad2deg(alpha);
    abc_angles[4] = rad2deg(beta);
    abc_angles[5] = rad2deg(gamma);
}
void abc_angles_to_vec(const double* abc_angles, double* vec)
{
    const double a = abc_angles[0];
    const double b = abc_angles[1];
    const double c = abc_angles[2];
    const double alpha = abc_angles[3];
    const double beta = abc_angles[4];
    const double gamma = abc_angles[5];
    vec[0] = a;
    vec[1] = 0.0;
    vec[2] = 0.0;
    vec[3] = b * std::cos(deg2rad(gamma));
    vec[4] = b * std::sin(deg2rad(gamma));
    vec[5] = 0.0;
    vec[6] = c * std::cos(deg2rad(beta));
    vec[7] = c * (std::cos(deg2rad(alpha)) - std::cos(deg2rad(beta)) * std::cos(deg2rad(gamma))) / std::sin(deg2rad(gamma));
    vec[8] = std::sqrt(c * c - vec[6] * vec[6] - vec[7] * vec[7]);
}
double vec_to_volume(const double* vec)
{
    // vector's mixed product
    return vec[0] * (vec[4] * vec[8] - vec[5] * vec[7]) - vec[1] * (vec[3] * vec[8] - vec[5] * vec[6]) + vec[2] * (vec[3] * vec[7] - vec[4] * vec[6]);
}
double abc_angles_to_volume(const double* abc_angles)
{
    std::vector vec(9);
    abc_angles_to_vec(abc_angles, vec.data());
    return vec_to_volume(vec.data());
}
std::vector _split_outside_enclose(const std::string& in, 
                                                const std::string& delim,
                                                const std::vector& enclose)
{
    // a very naive impl. for only CIF possible cases
    assert(enclose.size() == 2); // other complicated case not implemented yet
    // first split with delim. then scan all fragments, if there are enclose symbol, then will first meet
    // a fragment startswith the opening, then after fragments, there will be a one ends with closing.
    // between them, fragments will be concatenated with delim.
    std::vector out;
    std::string cache;
    std::vector words = FmtCore::split(in, delim);
    bool in_enclose = false;
    for (auto& word: words)
    {
        if (FmtCore::startswith(word, enclose[0]))
        {
            in_enclose = true;
            cache += word + delim;
        }
        else if (FmtCore::endswith(word, enclose[1]))
        {
            in_enclose = false;
            cache += word;
            out.push_back(cache);
            cache.clear();
        }
        else
        {
            if (in_enclose)
            {
                cache += word + delim;
            }
            else
            {
                out.push_back(word);
            }
        }
    }
    return out;
}

// the second step, for each block, split with words starting with "_"
std::vector _split_loop_block(const std::string& block)
{
    std::vector out;
    std::string word, cache;
    std::stringstream ss(FmtCore::strip(FmtCore::strip(block, "\n")));
    while (ss.good())
    {
        ss >> word;
        if (FmtCore::startswith(word, "_"))
        {
            if (!cache.empty())
            {
                out.push_back(FmtCore::strip(cache));
                cache.clear();
            }
            out.push_back(FmtCore::strip(word));
        }
        else
        {
            cache += word + " ";
        }
    }
    // the last word
    if (!cache.empty())
    {
        out.push_back(FmtCore::strip(cache));
    }
    return out;
}

std::map _build_table(const std::vector& keys,
                                                             const std::vector& values)
{
    std::map out;
    const size_t ncols = keys.size();
    assert(values.size() % ncols == 0);
    const size_t nrows = values.size() / ncols;
    for (size_t i = 0; i < ncols; i++)
    {
        std::vector col(nrows);
        for (size_t j = 0; j < nrows; j++)
        {
            col[j] = values[j * ncols + i];
        }
        out[keys[i]] = col;
    }
    return out;
}

std::map _build_block_data(const std::vector& block)
{
    // after calling the _split_loop_block, the data now composed of elements that either startswith "_"
    // or not. Between elements startswith "_", there is at most one element that does not startswith "_".
    // a scan can be performed to group those keys.
    std::vector keys;
    std::vector kcache;
    std::vector values;
    // first drop all elements that does not startswith "_" before the first element that startswith "_"
    std::vector block_ = block;
    auto it = std::find_if(block.begin(), block.end(), [](const std::string& s) { return FmtCore::startswith(s, "_"); });
    if (it != block.begin())
    {
        block_.erase(block_.begin(), it);
    }
    for (auto& elem: block_)
    {
        if (FmtCore::startswith(elem, "_"))
        {
            kcache.push_back(elem);
        }
        else
        {
            keys.push_back(kcache);
            values.push_back(elem);
            kcache.clear();
        }
    }
    assert(keys.size() == values.size()); // ensure the number of keys and values are the same
    // then for each elem in keys, if there are more than one element, then it is a table. Make it a table
    // , otherwise it is a simple key-value pair, directly add it to the output.
    std::map out;
    for (size_t i = 0; i < keys.size(); i++)
    {
        if (keys[i].size() > 1)
        {
            const std::vector words = _split_outside_enclose(values[i], " ", {"'", "'"});
            std::map table = _build_table(keys[i], words);
            out.insert(table.begin(), table.end());
        }
        else
        {
            out[keys[i][0]] = {values[i]};
        }
    }
    
    return out;
}

void bcast_cifmap(std::map& map, // the map to be broadcasted 
                  const int rank = 0)   // source rank: from which rank to broadcast
{
#ifdef __MPI
    int myrank;
    MPI_Comm_rank(MPI_COMM_WORLD, &myrank);
    // use Parallel_Common namespace bcast_int and bcast_string to broadcast the size of map and key, value pairs
    int size = map.size();
    Parallel_Common::bcast_int(size); // seems it can only broadcast from rank 0, so presently no need to specify
                                      // the rank to broadcast
    std::vector keys(size);
    std::vector values(size);
    int i = 0;
    if (myrank == rank) // if the rank is the source rank, then pack the map to keys and values
    {
        for (auto& elem: map)
        {
            keys[i] = elem.first;
            values[i] = elem.second;
            i++;
        }
    }
    for (int i = 0; i < size; i++)
    {
        Parallel_Common::bcast_string(keys[i]);
        int valsize = values[i].size();
        Parallel_Common::bcast_int(valsize);
        values[i].resize(valsize);
        Parallel_Common::bcast_string(values[i].data(), valsize);
    }
    if (myrank != rank) // if the rank is not the source rank, then unpack the keys and values to map
    {
        map.clear();
        for (int i = 0; i < size; i++)
        {
            map[keys[i]] = values[i];
        }
    }
#endif
}

void ModuleIO::CifParser::_unpack_ucell(const UnitCell& ucell,
                                        std::vector& veca,
                                        std::vector& vecb,
                                        std::vector& vecc,
                                        int& natom,
                                        std::vector& atom_site_labels,
                                        std::vector& atom_site_fract_coords)
{
    const double bohr2angstrom = 0.52917721067;
    const double lat0 = ucell.lat.lat0;
    veca.resize(3);
    vecb.resize(3);
    vecc.resize(3);
    veca[0] = ucell.a1.x;
    veca[1] = ucell.a1.y;
    veca[2] = ucell.a1.z;
    vecb[0] = ucell.a2.x;
    vecb[1] = ucell.a2.y;
    vecb[2] = ucell.a2.z;
    vecc[0] = ucell.a3.x;
    vecc[1] = ucell.a3.y;
    vecc[2] = ucell.a3.z;
    std::for_each(veca.begin(), veca.end(), [lat0, bohr2angstrom](double& x) { x *= lat0 * bohr2angstrom; });
    std::for_each(vecb.begin(), vecb.end(), [lat0, bohr2angstrom](double& x) { x *= lat0 * bohr2angstrom; });
    std::for_each(vecc.begin(), vecc.end(), [lat0, bohr2angstrom](double& x) { x *= lat0 * bohr2angstrom; });
    natom = ucell.nat;
    assert(natom > 0); // ensure the number of atoms is positive
    atom_site_labels.resize(natom);
    atom_site_fract_coords.resize(3 * natom);
    for (int i = 0; i < natom; ++i)
    {
        atom_site_labels[i] = ucell.atoms[ucell.iat2it[i]].ncpp.psd; // the most standard label
        atom_site_labels[i] = atom_site_labels[i].empty() ? ucell.atom_label[ucell.iat2it[i]]: atom_site_labels[i];
        atom_site_labels[i] = atom_site_labels[i].empty() ? ucell.atoms[ucell.iat2it[i]].label: atom_site_labels[i];
        assert(!atom_site_labels[i].empty()); // ensure the label is not empty
        atom_site_fract_coords[3 * i] = ucell.atoms[ucell.iat2it[i]].taud[ucell.iat2ia[i]].x;
        atom_site_fract_coords[3 * i + 1] = ucell.atoms[ucell.iat2it[i]].taud[ucell.iat2ia[i]].y;
        atom_site_fract_coords[3 * i + 2] = ucell.atoms[ucell.iat2it[i]].taud[ucell.iat2ia[i]].z;
    }
}

void ModuleIO::CifParser::write(const std::string& fcif,
                                const double* abc_angles,
                                const int natom,
                                const std::string* atom_site_labels,
                                const double* atom_site_fract_coords,
                                const std::string& title,
                                const std::string& data_tag,
                                const int rank,
                                const double* atom_site_occups,
                                const std::string& cell_formula_units_z)
{
#ifdef __MPI // well...very simple...
    int myrank;
    MPI_Comm_rank(MPI_COMM_WORLD, &myrank);
    if (myrank != rank) // if present rank is not the rank assigned to write the cif file, then return
    {
        return;
    }
#endif
    std::ofstream ofs(fcif);
    if (!ofs)
    {
        ModuleBase::WARNING_QUIT("ModuleIO::CifParser::write", "Cannot open file " + fcif);
    }
    ofs