import { Grid } from "@mui/material";
import React from "react";
import Plot from "react-plotly.js";

const colors = [
  "red",
  "blue",
  "green",
  "orange",
  "purple",
  "brown",
  "pink",
  "grey",
  "yellow",
  "cyan",
];

export default function EmbeddingPlot({ words, x, y, z }) {
  console.log(words);
  return (
    <>
      <Grid item xs={12}>
        <Plot
          data={[
            ...words.map((word, index) => {
              const color = colors[index % colors.length];
              return {
                x: [0, word.embedding[x]],
                y: [0, word.embedding[y]],
                z: [0, word.embedding[z]],
                type: "scatter3d",
                mode: "lines+markers",
                line: {
                  width: 6,
                },
                marker: {
                  size: 1,
                  color: color,
                },
                name: word.item,
              };
            }),
            ...words.map((word, index) => {
              const color = colors[index % colors.length];
              return {
                type: "cone",
                x: [word.embedding[x]],
                y: [word.embedding[y]],
                z: [word.embedding[z]],
                u: [word.embedding[x]],
                v: [word.embedding[y]],
                w: [word.embedding[z]],
                sizemode: "absolute",
                sizeref: 0.05,
                showscale: false,
                colorscale: [
                  [0, color],
                  [1, color],
                ],
              };
            }),
          ]}
          layout={{
            title: "3D Plot of Vectors",
            scene: {
              xaxis: { title: "X Axis" },
              yaxis: { title: "Y Axis" },
              zaxis: { title: "Z Axis" },
            },
            autosize: true,
          }}
          style={{ width: "100%", height: "100%" }}
        />
      </Grid>
    </>
  );
}