package us.ihmc.robotEnvironmentAwareness.ui.graphicsBuilders;

import java.util.concurrent.atomic.AtomicReference;

import perception_msgs.msg.dds.PlanarRegionsListMessage;
import javafx.scene.paint.Color;
import javafx.scene.paint.Material;
import javafx.scene.shape.Mesh;
import javafx.util.Pair;
import us.ihmc.communication.packets.PlanarRegionMessageConverter;
import us.ihmc.euclid.geometry.ConvexPolygon2D;
import us.ihmc.euclid.transform.RigidBodyTransform;
import us.ihmc.javaFXToolkit.shapes.JavaFXMultiColorMeshBuilder;
import us.ihmc.javaFXToolkit.shapes.TextureColorPalette2D;
import us.ihmc.messager.MessagerAPIFactory.Topic;
import us.ihmc.robotEnvironmentAwareness.communication.REAModuleAPI;
import us.ihmc.robotEnvironmentAwareness.communication.REAUIMessager;
import us.ihmc.robotics.geometry.PlanarRegion;
import us.ihmc.robotics.geometry.PlanarRegionsList;

public class PlanarRegionsMeshBuilder implements Runnable
{
   private final AtomicReference<Boolean> enable;
   private final AtomicReference<Boolean> clear;
   private final AtomicReference<Boolean> clearOcTree;

   private final JavaFXMultiColorMeshBuilder meshBuilder;

   private final AtomicReference<PlanarRegionsListMessage> planarRegionsListMessage;

   private final AtomicReference<Pair<Mesh, Material>> meshAndMaterialToRender = new AtomicReference<>(null);

   private final Topic<Boolean> requestPlanarRegionsTopic;
   
   private final REAUIMessager uiMessager;

   public PlanarRegionsMeshBuilder(REAUIMessager uiMessager, Topic<PlanarRegionsListMessage> planarRegionsState)
   {
      this(uiMessager, planarRegionsState, REAModuleAPI.OcTreeEnable, REAModuleAPI.PlanarRegionsPolygonizerClear, REAModuleAPI.OcTreeClear, REAModuleAPI.RequestPlanarRegions);
   }

   public PlanarRegionsMeshBuilder(REAUIMessager uiMessager, Topic<PlanarRegionsListMessage> planarRegionsState, Topic<Boolean> enableTopic,
                                   Topic<Boolean> clearTopic, Topic<Boolean> octreeClearTopic, Topic<Boolean> requestPlanarRegionsTopic)
   {
      enable = uiMessager.createInput(enableTopic, false);
      this.uiMessager = uiMessager;
      clear = uiMessager.createInput(clearTopic, false);
      clearOcTree = uiMessager.createInput(octreeClearTopic, false);
      this.requestPlanarRegionsTopic = requestPlanarRegionsTopic;

      planarRegionsListMessage = uiMessager.createInput(planarRegionsState);

      TextureColorPalette2D colorPalette = new TextureColorPalette2D();
      colorPalette.setHueBrightnessBased(0.9);
      meshBuilder = new JavaFXMultiColorMeshBuilder(colorPalette);
   }

   public void clear()
   {
      meshAndMaterialToRender.set(new Pair<>(null, null));
   }

   @Override
   public void run()
   {
      PlanarRegionsListMessage newMessage = planarRegionsListMessage.getAndSet(null);

      // Reset both clears by using only one pipe
      if (clearOcTree.getAndSet(false) | clear.getAndSet(false))
      {
         meshAndMaterialToRender.set(new Pair<>(null, null));
         return;
      }

      if (!enable.get())
         return;

      uiMessager.submitStateRequestToModule(requestPlanarRegionsTopic);

      if (newMessage == null || newMessage.getRegionId().size() == 0)
         return;

      meshAndMaterialToRender.set(generateMeshAndMaterial(newMessage));

   }

   private Pair<Mesh, Material> generateMeshAndMaterial(PlanarRegionsListMessage newMessage)
   {
      meshBuilder.clear();

      double lineWidth = 0.01;
      RigidBodyTransform transformToWorld = new RigidBodyTransform();
      PlanarRegionsList planarRegionsList = PlanarRegionMessageConverter.convertToPlanarRegionsList(newMessage);

      for (int regionIndex = 0; regionIndex < planarRegionsList.getNumberOfPlanarRegions(); regionIndex++)
      {
         PlanarRegion planarRegion = planarRegionsList.getPlanarRegion(regionIndex);

         int regionId = planarRegion.getRegionId();
         Color regionColor = getRegionColor(regionId);
         planarRegion.getTransformToWorld(transformToWorld);

         meshBuilder.addMultiLine(transformToWorld, planarRegion.getConcaveHull(), lineWidth, regionColor, true);

         for (int polygonIndex = 0; polygonIndex < planarRegion.getNumberOfConvexPolygons(); polygonIndex++)
         {
            ConvexPolygon2D convexPolygon2d = planarRegion.getConvexPolygon(polygonIndex);
            regionColor = Color.hsb(regionColor.getHue(), 0.9, 0.5 + 0.5 * ((double) polygonIndex / (double) planarRegion.getNumberOfConvexPolygons()));
            meshBuilder.addPolygon(transformToWorld, convexPolygon2d, regionColor);
         }
      }

      Material material = meshBuilder.generateMaterial();
      Mesh mesh = meshBuilder.generateMesh();

      return new Pair<>(mesh, material);
   }

   private Color getRegionColor(int regionId)
   {
      java.awt.Color awtColor = new java.awt.Color(regionId);
      return Color.rgb(awtColor.getRed(), awtColor.getGreen(), awtColor.getBlue());
   }

   public boolean hasNewMeshAndMaterial()
   {
      return meshAndMaterialToRender.get() != null;
   }

   public Pair<Mesh, Material> pollMeshAndMaterial()
   {
      return meshAndMaterialToRender.getAndSet(null);
   }
}
