#include "boundarylayer.hpp"

namespace netgen
{

namespace detail
{
struct Neighbor
{
  PointIndex pi;
  SurfaceElementIndex sei;
  double weight;
};
} // namespace detail

Array<ArrayMem<detail::Neighbor, 20>>
BuildNeighbors (FlatArray<PointIndex> points, const Mesh& mesh)
{
  auto p2sel = mesh.CreatePoint2SurfaceElementTable();

  Array<ArrayMem<detail::Neighbor, 20>> neighbors(points.Size());

  ArrayMem<double, 20> angles;
  ArrayMem<double, 20> inv_dists;
  for (auto i : points.Range())
    {
      auto& p_neighbors = neighbors[i];
      auto pi = points[i];
      angles.SetSize(0);
      inv_dists.SetSize(0);
      for (auto sei : p2sel[pi])
        {
          const auto& sel = mesh[sei];
          for (auto pi1 : sel.PNums())
            {
              if (pi1 == pi)
                continue;
              auto pi2 = pi1;
              for (auto pi_ : sel.PNums())
                {
                  if (pi_ != pi && pi_ != pi1)
                    {
                      pi2 = pi_;
                      break;
                    }
                }
              p_neighbors.Append({pi1, sei, 0.0});
              inv_dists.Append(1.0 / (mesh[pi1] - mesh[pi]).Length());
              auto dot = (mesh[pi1] - mesh[pi]).Normalize() * (mesh[pi2] - mesh[pi]).Normalize();
              angles.Append(acos(dot));
            }
        }
      double sum_inv_dist = 0.0;
      for (auto inv_dist : inv_dists)
        sum_inv_dist += inv_dist;
      double sum_angle = 0.0;
      for (auto angle : angles)
        sum_angle += angle;

      double sum_weight = 0.0;
      for (auto i : Range(inv_dists))
        {
          p_neighbors[i].weight =
            inv_dists[i] * angles[i] / sum_inv_dist / sum_angle;
          sum_weight += p_neighbors[i].weight;
        }
      for (auto i : Range(inv_dists))
        p_neighbors[i].weight /= sum_weight;
    }
  return neighbors;
}

void BoundaryLayerTool ::InterpolateGrowthVectors()
{
  int new_max_edge_nr = max_edge_nr;
  for (const auto& seg : segments)
    if (seg.edgenr > new_max_edge_nr)
      new_max_edge_nr = seg.edgenr;
  for (const auto& seg : new_segments)
    if (seg.edgenr > new_max_edge_nr)
      new_max_edge_nr = seg.edgenr;

  auto getGW = [&] (PointIndex pi) -> Vec<3> {
    if (growth_vector_map.count(pi) == 0)
      growth_vector_map[pi] = {&growthvectors[pi], total_height};
    auto [gw, height] = growth_vector_map[pi];
    return height * (*gw);
  };
  auto addGW = [&] (PointIndex pi, Vec<3> vec) {
    if (growth_vector_map.count(pi) == 0)
      growth_vector_map[pi] = {&growthvectors[pi], total_height};
    auto [gw, height] = growth_vector_map[pi];
    *gw += 1.0 / height * vec;
  };

  // interpolate tangential component of growth vector along edge
  if (max_edge_nr >= new_max_edge_nr)
    return;

  auto edgenr2seg = ngcore::CreateSortedTable<Segment*, int>(
    Range(segments.Size() + new_segments.Size()),
    [&] (auto& table, size_t segi) {
      auto& seg = segi < segments.Size()
                    ? segments[segi]
                    : new_segments[segi - segments.Size()];
      table.Add(seg.edgenr, &seg);
    },
    new_max_edge_nr + 1);
  auto point2seg = ngcore::CreateSortedTable<Segment*, PointIndex>(
    Range(segments.Size() + new_segments.Size()),
    [&] (auto& table, size_t segi) {
      auto& seg = segi < segments.Size()
                    ? segments[segi]
                    : new_segments[segi - segments.Size()];
      table.Add(seg[0], &seg);
      table.Add(seg[1], &seg);
    },
    mesh.GetNP());

  for (auto edgenr : Range(max_edge_nr + 1, new_max_edge_nr + 1))
    {
      double edge_len = 0.;

      auto is_end_point = [&] (PointIndex pi) {
        auto segs = point2seg[pi];
        if (segs.Size() == 1)
          return true;
        auto first_edgenr = (*segs[0]).edgenr;
        for (auto* p_seg : segs)
          if (p_seg->edgenr != first_edgenr)
            return true;
        return false;
      };

      bool any_grows = false;

      Array<PointIndex> points;
      for (auto* p_seg : edgenr2seg[edgenr])
        {
          auto& seg = *p_seg;

          if (getGW(seg[0]).Length2() != 0 || getGW(seg[1]).Length2() != 0)
            any_grows = true;

          if (points.Size() == 0)
            for (auto i : Range(2))
              if (is_end_point(seg[i]))
                {
                  points.Append(seg[i]);
                  points.Append(seg[1 - i]);
                  edge_len += (mesh[seg[1]] - mesh[seg[0]]).Length();
                  break;
                }
        }

      if (!any_grows)
        {
          PrintMessage(1, "BLayer: skip interpolating growth vectors at edge ", edgenr + 1);
          continue;
        }

      if (!points.Size())
        {
          cerr << "Could not find startpoint for edge " << edgenr << endl;
          continue;
        }

      std::set<PointIndex> points_set;
      points_set.insert(points[0]);
      points_set.insert(points[1]);

      bool point_found = true;
      while (point_found)
        {
          if (is_end_point(points.Last()))
            break;
          point_found = false;
          for (auto* p_seg : point2seg[points.Last()])
            {
              const auto& seg = *p_seg;
              if (seg.edgenr != edgenr)
                continue;
              auto plast = points.Last();
              if (plast != seg[0] && plast != seg[1])
                continue;
              auto pnew = plast == seg[0] ? seg[1] : seg[0];
              if (pnew == points[0] && points.Size() > 1)
                {
                }
              if (points_set.count(pnew) > 0 && (pnew != points[0] || points.Size() == 2))
                continue;
              edge_len += (mesh[points.Last()] - mesh[pnew]).Length();
              points.Append(pnew);
              points_set.insert(pnew);
              point_found = true;
              break;
            }
        }
      if (!point_found)
        {
          cerr << "Could not find connected list of line segments for edge "
               << edgenr << endl;
          cerr << "current points: " << endl
               << points << endl;
          continue;
        }

      if (getGW(points[0]).Length2() == 0 && getGW(points.Last()).Length2() == 0)
        continue;

      // tangential part of growth vectors
      auto t1 = (mesh[points[1]] - mesh[points[0]]).Normalize();
      auto gt1 = getGW(points[0]) * t1 * t1;
      auto t2 =
        (mesh[points.Last()] - mesh[points[points.Size() - 2]]).Normalize();
      auto gt2 = getGW(points.Last()) * t2 * t2;

      double len = 0.;
      for (auto i : IntRange(1, points.Size() - 1))
        {
          auto pi = points[i];
          len += (mesh[pi] - mesh[points[i - 1]]).Length();
          auto t = getEdgeTangent(pi, edgenr, point2seg[pi]);
          auto lam = len / edge_len;
          auto interpol = (1 - lam) * (gt1 * t) * t + lam * (gt2 * t) * t;
          addGW(pi, interpol);
        }
    }
}

void BoundaryLayerTool ::InterpolateSurfaceGrowthVectors()
{
  static Timer tall("InterpolateSurfaceGrowthVectors");
  RegionTimer rtall(tall);
  static Timer tsmooth("InterpolateSurfaceGrowthVectors-Smoothing");
  auto np_old = this->np;
  auto np = mesh.GetNP();

  // non_bl_growth_vectors.clear();
  // for(const auto & sel : new_sels) {
  //   for(auto pi : sel.PNums()) {
  //     if(mesh[pi].Type() == INNERPOINT)
  //       mesh[pi].SetType(SURFACEPOINT);
  //   }
  // }
  // cout << __FILE__ << ":" << __LINE__ << endl;

  auto getGW = [&] (PointIndex pi) -> Vec<3> {
    return growthvectors[pi];
    // if (growth_vector_map.count(pi) == 0) {
    //   non_bl_growth_vectors[pi] = .0;
    //   growth_vector_map[pi] = {&non_bl_growth_vectors[pi], 1.0};
    // }
    // auto [gw, height] = growth_vector_map[pi];
    // return height * (*gw);
  };
  auto addGW = [&] (PointIndex pi, Vec<3> vec) {
    growthvectors[pi] += vec;
    // // cout << "add gw " << pi << "\t" << vec << endl;
    // if (growth_vector_map.count(pi) == 0) {
    //   // cout << "\t make new entry" << endl;
    //   non_bl_growth_vectors[pi] = .0;
    //   growth_vector_map[pi] = {&non_bl_growth_vectors[pi], 1.0};
    // }
    // auto [gw, height] = growth_vector_map[pi];
    // // cout << "\tcurrent gw " << *gw << "\t" << height << endl;
    // *gw += 1.0 / height * vec;
  };

  auto hasMoved = [&] (PointIndex pi) {
    return (pi - PointIndex::BASE >= np_old) || mapto[pi].Size() > 0 || special_boundary_points.count(pi);
  };

  // cout << __FILE__ << ":" << __LINE__ << endl;
  std::set<PointIndex> points_set;
  for (const auto& sel : mesh.SurfaceElements())
    {
      for (auto pi : sel.PNums())
        if (mesh[pi].Type() == SURFACEPOINT && hasMoved(pi))
          points_set.insert(pi);
    }

  Array<PointIndex> points;
  for (auto pi : points_set)
    points.Append(pi);
  QuickSort(points);

  // smooth tangential part of growth vectors from edges to surface elements
  Array<Vec<3>, PointIndex> corrections(mesh.GetNP());
  corrections = 0.0;
  RegionTimer rtsmooth(tsmooth);
  auto neighbors = BuildNeighbors(points, mesh);

  Array<Vec<3>, SurfaceElementIndex> surf_normals(mesh.GetNSE());
  for (auto sei : mesh.SurfaceElements().Range())
    surf_normals[sei] = getNormal(mesh[sei]);

  BitArray interpolate_tangent(mesh.GetNP() + 1);
  interpolate_tangent = false;
  for (auto pi : points)
    {
      for (auto sei : p2sel[pi])
        if (is_boundary_moved[mesh[sei].GetIndex()])
          interpolate_tangent.SetBit(pi);
    }

  constexpr int N_STEPS = 64;
  for ([[maybe_unused]] auto i : Range(N_STEPS))
    {
      for (auto i : points.Range())
        {
          auto pi = points[i];
          // cout << "AVERAGE " << pi << endl;
          auto& p_neighbors = neighbors[i];

          ArrayMem<Vec<3>, 20> g_vectors;
          double max_len = 0.0;
          double sum_len = 0.0;

          // average only tangent component on new bl points, average whole growth
          // vector otherwise
          bool do_average_tangent = true;
          for (const auto& s : p_neighbors)
            {
              auto gw_other = getGW(s.pi) + corrections[s.pi];
              if (do_average_tangent)
                {
                  auto n = surf_normals[s.sei];
                  gw_other = gw_other - (gw_other * n) * n;
                }
              auto v = gw_other;
              auto len = v.Length2();
              sum_len += len;
              max_len = max(max_len, len);
              g_vectors.Append(v);
            }

          if (max_len == 0.0)
            continue;

          double lambda = 0;
          if (i > N_STEPS / 4.)
            lambda = 2.0 * (i - N_STEPS / 4.) / (N_STEPS / 2.);
          lambda = min(1.0, lambda);

          auto& correction = corrections[pi];
          correction = 0.0;
          for (const auto i : p_neighbors.Range())
            {
              auto v = g_vectors[i];
              double weight = lambda * p_neighbors[i].weight + (1.0 - lambda) * v.Length2() / sum_len;
              // if(pi == 19911) cout << "pi " << pi << "\tneighbor " <<
              // p_neighbors[i].pi << "\tweight " << weight << endl;
              correction += weight * v;
            }

          if (!do_average_tangent)
            correction -= getGW(pi);
          // if(pi == 19911) cout << "pi " << pi << "\tcorrection " << correction <<
          // endl;
        }
    }

  for (auto pi : points)
    addGW(pi, corrections[pi]);
}

void BoundaryLayerTool ::FixSurfaceElements()
{
  static Timer tall("FixSurfaceElements");
  RegionTimer rtall(tall);
  auto np_old = this->np;
  auto np = mesh.GetNP();

  non_bl_growth_vectors.clear();

  auto getGW = [&] (PointIndex pi) -> Vec<3> {
    // return growthvectors[pi];
    if (growth_vector_map.count(pi) == 0)
      {
        non_bl_growth_vectors[pi] = .0;
        growth_vector_map[pi] = {&non_bl_growth_vectors[pi], 1.0};
      }
    auto [gw, height] = growth_vector_map[pi];
    return height * (*gw);
  };

  auto addGW = [&] (PointIndex pi, Vec<3> vec) {
    if (growth_vector_map.count(pi) == 0)
      {
        non_bl_growth_vectors[pi] = .0;
        growth_vector_map[pi] = {&non_bl_growth_vectors[pi], 1.0};
      }
    auto [gw, height] = growth_vector_map[pi];
    *gw += 1.0 / height * vec;
  };

  std::set<PointIndex> points_set;
  // only smooth over old surface elements
  for (SurfaceElementIndex sei : Range(nse))
    {
      const auto& sel = mesh[sei];
      if (sel.GetNP() == 3 && is_boundary_moved[sel.GetIndex()])
        for (auto pi : sel.PNums())
          if (mesh[pi].Type() == SURFACEPOINT)
            points_set.insert(pi);
    }

  Array<PointIndex> points;
  for (auto pi : points_set)
    points.Append(pi);
  QuickSort(points);

  Array<Vec<3>, PointIndex> corrections(mesh.GetNP());
  corrections = 0.0;

  auto neighbors = BuildNeighbors(points, mesh);

  constexpr int N_STEPS = 32;
  for ([[maybe_unused]] auto i : Range(N_STEPS))
    {
      for (auto i : points.Range())
        {
          auto pi = points[i];
          auto& p_neighbors = neighbors[i];

          ArrayMem<Vec<3>, 20> g_vectors;
          double max_len = 0.0;
          double sum_len = 0.0;

          for (const auto& s : p_neighbors)
            {
              auto v = getGW(s.pi) + corrections[s.pi];
              auto len = v.Length2();
              sum_len += len;
              max_len = max(max_len, len);
              g_vectors.Append(v);
            }

          if (max_len == 0.0)
            continue;

          double lambda = 0;
          if (i > N_STEPS / 4.)
            lambda = 2.0 * (i - N_STEPS / 4.) / (N_STEPS / 2.);
          lambda = min(1.0, lambda);

          auto& correction = corrections[pi];
          correction = 0.0;
          for (const auto i : p_neighbors.Range())
            {
              auto v = g_vectors[i];
              double weight = lambda * p_neighbors[i].weight + (1.0 - lambda) * v.Length2() / sum_len;
              correction += weight * v;
            }
        }
    }

  for (auto pi : points)
    addGW(pi, corrections[pi]);
}

} // namespace netgen