import React, { useState, useEffect } from "react";
import axios from "axios";
import { ResponsiveSankey } from "@nivo/sankey";
import { ResponsiveLine } from "@nivo/line";
import PointCloud from "../components/PointCloud";
import "../style.css";

const TemporalAnalysis: React.FC = () => {
  const [cellDataCS7, setCellDataCS7] = useState<any[]>([]);
  const [cellDataCS8, setCellDataCS8] = useState<any[]>([]);
  const [cellDataCS9, setCellDataCS9] = useState<any[]>([]);
  const [loadingCS7, setLoadingCS7] = useState(false);
  const [loadingCS8, setLoadingCS8] = useState(false);
  const [loadingCS9, setLoadingCS9] = useState(false);
  const [errorCS7, setErrorCS7] = useState("");
  const [errorCS8, setErrorCS8] = useState("");
  const [errorCS9, setErrorCS9] = useState("");
  const [sankeyData, setSankeyData] = useState<any>({ nodes: [], links: [] });
  const [sankeyLoading, setSankeyLoading] = useState(false);

  // 第三部分：基因时序分析相关状态
  const [availableGenes, setAvailableGenes] = useState<string[]>([]);
  const [selectedGene, setSelectedGene] = useState("");
  const [temporalData, setTemporalData] = useState<any[]>([]);
  const [temporalLoading, setTemporalLoading] = useState(false);
  const [temporalError, setTemporalError] = useState("");

  // 生成细胞类型颜色映射
  const generateCellTypeColors = (cellTypes: string[]) => {
    const colorMap = new Map<string, string>();
    const colors = [
      '#FF6B6B', '#4ECDC4', '#45B7D1', '#96CEB4', '#FFEAA7',
      '#DDA0DD', '#98D8C8', '#F7DC6F', '#BB8FCE', '#85C1E9',
      '#F8C471', '#82E0AA', '#F1948A', '#85CDCE', '#D7BDE2',
      '#A9DFBF', '#F9E79F', '#D5A6BD', '#AED6F1', '#E8DAEF'
    ];

    cellTypes.forEach((type, index) => {
      colorMap.set(type, colors[index % colors.length]);
    });

    return colorMap;
  };

  // 获取可用基因列表（从CS7阶段获取）
  useEffect(() => {
    axios
      .get("http://localhost:5000/api/genes", {
        params: { stage: "CS7" },
      })
      .then((res) => setAvailableGenes(res.data))
      .catch(() => setAvailableGenes([]));
  }, []);

  // 获取基因时序分析数据
  const fetchTemporalAnalysis = async (gene: string) => {
    if (!gene) return;

    setTemporalLoading(true);
    setTemporalError("");
    setTemporalData([]);

    try {
      const response = await axios.get("http://localhost:5000/api/gene_temporal_analysis", {
        params: { gene: gene.trim() }
      });

      // 转换数据为线图格式
      const stagesData = response.data.stages_data;
      const allCellTypes = new Set<string>();

      // 收集所有细胞类型
      stagesData.forEach((stageData: any) => {
        Object.keys(stageData.cell_types).forEach(cellType => {
          allCellTypes.add(cellType);
        });
      });

      const cellTypesArray = Array.from(allCellTypes).sort();
      const colorMap = generateCellTypeColors(cellTypesArray);

      // 构建线图数据 - 每个细胞类型一条线
      const lineData = cellTypesArray.map(cellType => {
        const linePoints = stagesData.map((stageData: any) => {
          const cellTypeData = stageData.cell_types[cellType];
          return {
            x: stageData.stage,
            y: cellTypeData ? cellTypeData.proportion : 0
          };
        });

        return {
          id: cellType,
          color: colorMap.get(cellType) || '#CCCCCC',
          data: linePoints
        };
      });

      setTemporalData(lineData);
      console.log('Temporal analysis data:', lineData);
    } catch (err: any) {
      if (err.response && err.response.data?.error) {
        setTemporalError(err.response.data.error);
      } else {
        setTemporalError("Failed to fetch temporal analysis data.");
      }
    } finally {
      setTemporalLoading(false);
    }
  };

  // 构建桑基图数据
  const buildSankeyData = (dataCS7: any[], dataCS8: any[], dataCS9: any[]) => {
    setSankeyLoading(true);

    try {
      // 统计各阶段的细胞类型数量
      const countCellTypes = (data: any[]) => {
        const counts: { [key: string]: number } = {};
        data.forEach(cell => {
          const cellType = cell.value || 'Unknown'; // 修正：细胞类型存储在value字段中
          counts[cellType] = (counts[cellType] || 0) + 1;
        });
        return counts;
      };

      const countsCS7 = countCellTypes(dataCS7);
      const countsCS8 = countCellTypes(dataCS8);
      const countsCS9 = countCellTypes(dataCS9);

      // 调试信息
      console.log('Cell type counts CS7:', countsCS7);
      console.log('Cell type counts CS8:', countsCS8);
      console.log('Cell type counts CS9:', countsCS9);

      // 获取所有细胞类型
      const allCellTypes = new Set([
        ...Object.keys(countsCS7),
        ...Object.keys(countsCS8),
        ...Object.keys(countsCS9)
      ]);

      const cellTypeArray = Array.from(allCellTypes).sort();
      const colorMap = generateCellTypeColors(cellTypeArray);

      // 计算总数
      const totalCS7 = Object.values(countsCS7).reduce((sum, count) => sum + count, 0);
      const totalCS8 = Object.values(countsCS8).reduce((sum, count) => sum + count, 0);
      const totalCS9 = Object.values(countsCS9).reduce((sum, count) => sum + count, 0);

      // 构建节点
      const nodes: any[] = [];
      cellTypeArray.forEach(cellType => {
        const color = colorMap.get(cellType) || '#CCCCCC';

        // CS7 节点
        if (countsCS7[cellType]) {
          const percentage = ((countsCS7[cellType] / totalCS7) * 100).toFixed(1);
          nodes.push({
            id: `CS7_${cellType}`,
            nodeColor: color,
            label: `${cellType} (${percentage}%)`
          });
        }

        // CS8 节点
        if (countsCS8[cellType]) {
          const percentage = ((countsCS8[cellType] / totalCS8) * 100).toFixed(1);
          nodes.push({
            id: `CS8_${cellType}`,
            nodeColor: color,
            label: `${cellType} (${percentage}%)`
          });
        }

        // CS9 节点
        if (countsCS9[cellType]) {
          const percentage = ((countsCS9[cellType] / totalCS9) * 100).toFixed(1);
          nodes.push({
            id: `CS9_${cellType}`,
            nodeColor: color,
            label: `${cellType} (${percentage}%)`
          });
        }
      });

      // 构建连接：使用最小百分比作为同类连线，并将剩余部分匹配到有缺口的目标细胞类型，保证各阶段总量守恒
      const buildBalancedLinks = (
        sourceCounts: Record<string, number>,
        sourceTotal: number,
        targetCounts: Record<string, number>,
        targetTotal: number,
        sourceStage: string,
        targetStage: string
      ) => {
        const types = Array.from(
          new Set([
            ...Object.keys(sourceCounts),
            ...Object.keys(targetCounts)
          ])
        );

        const sourcePct: Record<string, number> = {};
        const targetPct: Record<string, number> = {};
        types.forEach((t) => {
          sourcePct[t] = ((sourceCounts[t] || 0) / sourceTotal) * 100;
          targetPct[t] = ((targetCounts[t] || 0) / targetTotal) * 100;
        });

        const flows: any[] = [];

        // 同类对齐：取最小值
        const sourceResidual: Record<string, number> = {};
        const targetResidual: Record<string, number> = {};
        types.forEach((t) => {
          const direct = Math.min(sourcePct[t], targetPct[t]);
          if (direct > 0) {
            const color = (colorMap.get(t) || '#CCCCCC') + '80';
            flows.push({
              source: `${sourceStage}_${t}`,
              target: `${targetStage}_${t}`,
              value: direct,
              color
            });
          }
          sourceResidual[t] = sourcePct[t] - direct; // >= 0
          targetResidual[t] = targetPct[t] - direct; // >= 0
        });

        // 将源阶段的剩余按目标阶段的缺口比例分配
        types.forEach((sType) => {
          let remaining = sourceResidual[sType];
          if (remaining <= 0) return;
          // 计算当前目标缺口总和
          let totalDeficit = types.reduce((sum, t) => sum + Math.max(targetResidual[t], 0), 0);
          if (totalDeficit <= 0) return;
          types.forEach((dType) => {
            if (targetResidual[dType] <= 0) return;
            const share = (remaining * targetResidual[dType]) / totalDeficit;
            if (share > 1e-6) {
              const color = (colorMap.get(sType) || '#CCCCCC') + '80';
              flows.push({
                source: `${sourceStage}_${sType}`,
                target: `${targetStage}_${dType}`,
                value: share,
                color
              });
              targetResidual[dType] -= share;
            }
          });
        });

        return flows;
      };

      const links: any[] = [
        ...buildBalancedLinks(countsCS7, totalCS7, countsCS8, totalCS8, 'CS7', 'CS8'),
        ...buildBalancedLinks(countsCS8, totalCS8, countsCS9, totalCS9, 'CS8', 'CS9')
      ];

      setSankeyData({ nodes, links });
    } catch (error) {
      console.error('Error building Sankey data:', error);
      setSankeyData({ nodes: [], links: [] });
    } finally {
      setSankeyLoading(false);
    }
  };

  // 获取CS7阶段细胞类型数据
  useEffect(() => {
    setLoadingCS7(true);
    setErrorCS7("");
    setCellDataCS7([]);

    axios
      .get("http://localhost:5000/api/cell", {
        params: { stage: "CS7" },
      })
      .then((res) => {
        setCellDataCS7(res.data.cells);
      })
      .catch((err) => {
        if (err.response && err.response.data?.error) {
          setErrorCS7(err.response.data.error);
        } else {
          setErrorCS7("Failed to fetch CS7 cell type data.");
        }
      })
      .finally(() => {
        setLoadingCS7(false);
      });
  }, []);

  // 获取CS8阶段细胞类型数据
  useEffect(() => {
    setLoadingCS8(true);
    setErrorCS8("");
    setCellDataCS8([]);

    axios
      .get("http://localhost:5000/api/cell", {
        params: { stage: "CS8" },
      })
      .then((res) => {
        setCellDataCS8(res.data.cells);
      })
      .catch((err) => {
        if (err.response && err.response.data?.error) {
          setErrorCS8(err.response.data.error);
        } else {
          setErrorCS8("Failed to fetch CS8 cell type data.");
        }
      })
      .finally(() => {
        setLoadingCS8(false);
      });
  }, []);

  // 获取CS9阶段细胞类型数据
  useEffect(() => {
    setLoadingCS9(true);
    setErrorCS9("");
    setCellDataCS9([]);

    axios
      .get("http://localhost:5000/api/cell", {
        params: { stage: "CS9" },
      })
      .then((res) => {
        setCellDataCS9(res.data.cells);
      })
      .catch((err) => {
        if (err.response && err.response.data?.error) {
          setErrorCS9(err.response.data.error);
        } else {
          setErrorCS9("Failed to fetch CS9 cell type data.");
        }
      })
      .finally(() => {
        setLoadingCS9(false);
      });
  }, []);

  // 当所有数据加载完成后构建桑基图数据
  useEffect(() => {
    if (cellDataCS7.length > 0 && cellDataCS8.length > 0 && cellDataCS9.length > 0) {
      buildSankeyData(cellDataCS7, cellDataCS8, cellDataCS9);
    }
  }, [cellDataCS7, cellDataCS8, cellDataCS9]);

  return (
    <div className="page-container">
      <div className="page-header">
        <h1>Temporal Analysis</h1>
        <p className="page-description">
          Compare cell type distributions across stages to analyze changes and evolutionary patterns during embryonic development.
        </p>
      </div>

      <div className="page-content">
        {/* 第一部分：三个发育阶段的细胞类型分布 */}
        <div className="section">
          <h2>Stage comparison</h2>
          <p className="flat-section-description">
            The three 3D views below show cell type distributions for CS7, CS8, and CS9.
            Compare spatial distributions to observe changes over development.
          </p>

          <div
            className="three-column-grid visualization-grid"
            style={{
              display: "grid",
              gridTemplateColumns: "repeat(3, minmax(400px, 1fr))",
              gap: "2rem",
              minHeight: "600px",
              width: "100%",
              marginBottom: "3rem"
            }}
          >
            {/* CS7阶段细胞类型分布 */}
            <div className="visualization-container">
              <h3 style={{ marginBottom: "1rem", textAlign: "center" }}>CS7 - Early developmental stage</h3>
              {errorCS7 && <div className="error-message">{errorCS7}</div>}
              {cellDataCS7.length > 0 ? (
                <PointCloud data={cellDataCS7} isCategorical={true} />
              ) : (
                <div className="no-data-message">
                  {loadingCS7 ? (
                    <div>
                      <div style={{ fontSize: "2rem", marginBottom: "1rem" }}>🔄</div>
                       <p>Loading CS7 data...</p>
                    </div>
                  ) : (
                    <div>
                      <div style={{ fontSize: "2rem", marginBottom: "1rem" }}>🧬</div>
                       <p>CS7 cell type distribution</p>
                      <p style={{ fontSize: "0.9rem", color: "var(--text-secondary)" }}>
                         Early developmental stage
                      </p>
                    </div>
                  )}
                </div>
              )}
            </div>

            {/* CS8阶段细胞类型分布 */}
            <div className="visualization-container">
              <h3 style={{ marginBottom: "1rem", textAlign: "center" }}>CS8 - Mid developmental stage</h3>
              {errorCS8 && <div className="error-message">{errorCS8}</div>}
              {cellDataCS8.length > 0 ? (
                <PointCloud data={cellDataCS8} isCategorical={true} />
              ) : (
                <div className="no-data-message">
                  {loadingCS8 ? (
                    <div>
                      <div style={{ fontSize: "2rem", marginBottom: "1rem" }}>🔄</div>
                       <p>Loading CS8 data...</p>
                    </div>
                  ) : (
                    <div>
                      <div style={{ fontSize: "2rem", marginBottom: "1rem" }}>🧬</div>
                       <p>CS8 cell type distribution</p>
                      <p style={{ fontSize: "0.9rem", color: "var(--text-secondary)" }}>
                         Mid developmental stage
                      </p>
                    </div>
                  )}
                </div>
              )}
            </div>

            {/* CS9阶段细胞类型分布 */}
            <div className="visualization-container">
              <h3 style={{ marginBottom: "1rem", textAlign: "center" }}>CS9 - Late developmental stage</h3>
              {errorCS9 && <div className="error-message">{errorCS9}</div>}
              {cellDataCS9.length > 0 ? (
                <PointCloud data={cellDataCS9} isCategorical={true} />
              ) : (
                <div className="no-data-message">
                  {loadingCS9 ? (
                    <div>
                      <div style={{ fontSize: "2rem", marginBottom: "1rem" }}>🔄</div>
                       <p>Loading CS9 data...</p>
                    </div>
                  ) : (
                    <div>
                      <div style={{ fontSize: "2rem", marginBottom: "1rem" }}>🧬</div>
                       <p>CS9 cell type distribution</p>
                      <p style={{ fontSize: "0.9rem", color: "var(--text-secondary)" }}>
                         Late developmental stage
                      </p>
                    </div>
                  )}
                </div>
              )}
            </div>
          </div>
        </div>

                {/* 第二部分和第三部分：并排显示 */}
        <div className="section">
          <h2>Flow and gene timeline</h2>
          <p className="flat-section-description">
            The Sankey diagram (left) shows cell type flows across stages; the line chart (right)
            shows the proportion changes of the selected gene across cell types.
          </p>

          <div
            style={{
              display: "grid",
              gridTemplateColumns: "1fr 1fr",
              gap: "2rem",
              marginBottom: "3rem"
            }}
          >
            {/* 桑基图部分 */}
            <div className="visualization-container">
              <h3 style={{ marginBottom: "1rem", textAlign: "center" }}>Cell type flow diagram</h3>
              <div style={{ width: "100%", height: "500px", position: "relative" }}>
                {sankeyData.nodes.length > 0 && sankeyData.links.length > 0 ? (
                  <>
                    <ResponsiveSankey
                      data={sankeyData}
                      margin={{ top: 30, right: 100, bottom: 50, left: 100 }}
                      align="justify"
                      colors={(node: any) => node.nodeColor}
                      nodeOpacity={0.8}
                      nodeHoverOpacity={1}
                      nodeThickness={18}
                      nodeSpacing={8}
                      nodeBorderWidth={2}
                      nodeBorderColor={{
                        from: 'color',
                        modifiers: [['darker', 0.3]]
                      }}
                      linkOpacity={0.3}
                      linkHoverOpacity={0.6}
                      linkContract={4}
                      enableLinkGradient={true}
                      labelPosition="outside"
                      labelOrientation="horizontal"
                      labelPadding={12}
                      labelTextColor={{
                        from: 'color',
                        modifiers: [['darker', 1]]
                      }}
                      legends={[]}
                      motionConfig="gentle"
                    />
                      {/* Stage labels */}
                    <div
                      style={{
                        position: "absolute",
                        bottom: "15px",
                        left: "100px",
                        right: "100px",
                        display: "flex",
                        justifyContent: "space-between",
                        alignItems: "center",
                        fontSize: "13px",
                        fontWeight: "500",
                        color: "#666"
                      }}
                    >
                      <div style={{ textAlign: "left", flex: 1, paddingLeft: "20px" }}>
                        CS7 - Early
                      </div>
                      <div style={{ textAlign: "center", flex: 1 }}>
                        CS8 - Mid
                      </div>
                      <div style={{ textAlign: "right", flex: 1, paddingRight: "20px" }}>
                        CS9 - Late
                      </div>
                    </div>
                  </>
                ) : (
                  <div className="no-data-message">
                    {sankeyLoading || loadingCS7 || loadingCS8 || loadingCS9 ? (
                      <div>
                        <div style={{ fontSize: "2rem", marginBottom: "1rem" }}>🔄</div>
                        <p>Building flow diagram...</p>
                      </div>
                    ) : (
                      <div>
                        <div style={{ fontSize: "2rem", marginBottom: "1rem" }}>📊</div>
                        <p>Waiting for data</p>
                      </div>
                    )}
                  </div>
                )}
              </div>
            </div>

            {/* 基因时序分析部分 */}
            <div className="visualization-container">
              <h3 style={{ marginBottom: "1rem", textAlign: "center" }}>Gene temporal analysis</h3>
              <div style={{ display: "flex", justifyContent: "center", alignItems: "center", marginBottom: "1rem" }}>
                <select
                  value={selectedGene}
                  onChange={(e) => {
                    setSelectedGene(e.target.value);
                    if (e.target.value) {
                      fetchTemporalAnalysis(e.target.value);
                    }
                  }}
                  title="Select a gene for temporal analysis"
                  style={{ padding: "0.5rem", borderRadius: "0.3rem", border: "1px solid #ccc", fontSize: "0.9rem" }}
                >
                  <option value="">Please select a gene</option>
                  {availableGenes.map((gene) => (
                    <option key={gene} value={gene}>
                      {gene}
                    </option>
                  ))}
                </select>
              </div>

              <div style={{ width: "100%", height: "450px" }}>
                {temporalLoading ? (
                  <div className="no-data-message">
                    <div style={{ fontSize: "2rem", marginBottom: "1rem" }}>🔄</div>
                    <p>Loading temporal data...</p>
                  </div>
                ) : temporalError ? (
                  <div className="error-message">{temporalError}</div>
                ) : temporalData.length > 0 ? (
                  <ResponsiveLine
                    data={temporalData}
                    margin={{ top: 30, right: 70, bottom: 70, left: 70 }}
                    xScale={{ type: 'point' }}
                    yScale={{
                      type: 'linear',
                      min: 0,
                      max: 'auto',
                      stacked: false,
                      reverse: false
                    }}
                    yFormat=" >-.2f"
                    axisTop={null}
                    axisRight={null}
                      axisBottom={{
                      tickSize: 5,
                      tickPadding: 5,
                      tickRotation: 0,
                        legend: 'Developmental stage',
                      legendPosition: 'middle',
                      legendOffset: 40
                    }}
                      axisLeft={{
                      tickSize: 5,
                      tickPadding: 5,
                      tickRotation: 0,
                        legend: 'Proportion (%)',
                      legendPosition: 'middle',
                      legendOffset: -50
                    }}
                    colors={(line: any) => line.color}
                    pointSize={8}
                    pointColor={{ theme: 'background' }}
                    pointBorderWidth={3}
                    pointBorderColor={{ from: 'serieColor' }}
                    pointLabelYOffset={-12}
                    useMesh={true}
                    enableGridX={false}
                    enableGridY={true}
                    lineWidth={3}
                    legends={[
                      {
                        anchor: 'top',
                        direction: 'row',
                        justify: false,
                        translateX: 0,
                        translateY: -30,
                        itemsSpacing: 45,
                        itemWidth: 70,
                        itemHeight: 16,
                        itemDirection: 'left-to-right',
                        itemOpacity: 0.75,
                        symbolSize: 8,
                        symbolShape: 'circle',
                        effects: [
                          {
                            on: 'hover',
                            style: {
                              itemOpacity: 1
                            }
                          }
                        ]
                      }
                    ]}
                  />
                ) : (
                    <div className="no-data-message">
                    <div style={{ fontSize: "2rem", marginBottom: "1rem" }}>📊</div>
                    <p>Please select a gene to view temporal changes</p>
                  </div>
                )}
              </div>
            </div>
          </div>
        </div>

                 <div className="section">
           <h2>User Guide</h2>
           <div style={{ display: "grid", gridTemplateColumns: "repeat(auto-fit, minmax(300px, 1fr))", gap: "2rem", marginTop: "1.5rem" }}>
             <div>
               <h3>📅 Time-series comparison</h3>
               <p>The three 3D views for CS7, CS8, and CS9 allow intuitive observation of how cell types change over time.</p>
             </div>
             <div>
               <h3>🌊 Interpreting the Sankey diagram</h3>
               <p>Node size reflects proportions; links show continuity of the same cell type across stages. Colors are consistent to aid tracking.</p>
             </div>
             <div>
               <h3>📊 Gene temporal analysis</h3>
               <p>After selecting a gene, the line chart shows how its proportion changes across cell types, revealing expression dynamics during development.</p>
             </div>
             <div>
               <h3>🖱️ Interactions</h3>
               <p>Rotate with drag, zoom with scroll, click data points for details. Hover over the Sankey and line charts to see exact values.</p>
             </div>
             <div>
               <h3>🎨 Color encoding</h3>
               <p>Identical colors denote the same cell type across charts to make it easy to track changes and flows.</p>
             </div>
             <div>
               <h3>🔍 Developmental pattern identification</h3>
               <p>Compare stages and genes to identify emergence, disappearance, migration, and differentiation of cell types and their expression dynamics.</p>
             </div>
             <div>
               <h3>📈 Percentage analysis</h3>
               <p>Percentages indicate relative abundance of cell types per stage, helping interpret dynamic composition changes.</p>
             </div>
             <div>
               <h3>🧬 Gene selection strategy</h3>
               <p>Select development-related genes (e.g., SOX2, NANOG) to observe expression changes and cell-type specificity during embryogenesis.</p>
             </div>
           </div>
         </div>
      </div>
    </div>
  );
};

export default TemporalAnalysis;