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 links: any[] = [];
      cellTypeArray.forEach(cellType => {
        const color = colorMap.get(cellType) || '#CCCCCC';
        
        // CS7 -> CS8 连接
        if (countsCS7[cellType] && countsCS8[cellType]) {
          const cs7Percentage = (countsCS7[cellType] / totalCS7) * 100;
          const cs8Percentage = (countsCS8[cellType] / totalCS8) * 100;
          const avgPercentage = (cs7Percentage + cs8Percentage) / 2;
          
          links.push({
            source: `CS7_${cellType}`,
            target: `CS8_${cellType}`,
            value: Math.max(avgPercentage, 0.1), // 确保有最小值以显示连接
            color: color + '80' // 添加透明度
          });
        }

        // CS8 -> CS9 连接
        if (countsCS8[cellType] && countsCS9[cellType]) {
          const cs8Percentage = (countsCS8[cellType] / totalCS8) * 100;
          const cs9Percentage = (countsCS9[cellType] / totalCS9) * 100;
          const avgPercentage = (cs8Percentage + cs9Percentage) / 2;
          
          links.push({
            source: `CS8_${cellType}`,
            target: `CS9_${cellType}`,
            value: Math.max(avgPercentage, 0.1), // 确保有最小值以显示连接
            color: color + '80' // 添加透明度
          });
        }
      });

      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>时间序列分析</h1>
        <p className="page-description">
          通过对比不同发育阶段的细胞类型分布，分析胚胎发育过程中细胞类型的变化和演进模式。
        </p>
      </div>

      <div className="page-content">
        {/* 第一部分：三个发育阶段的细胞类型分布 */}
        <div className="section">
          <h2>发育阶段对比</h2>
          <p className="flat-section-description">
            下方展示了CS7、CS8、CS9三个发育阶段的细胞类型3D分布图。通过对比可以观察细胞类型在发育过程中的空间分布变化。
          </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 - 早期发育阶段</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>正在加载CS7数据...</p>
                    </div>
                  ) : (
                    <div>
                      <div style={{ fontSize: "2rem", marginBottom: "1rem" }}>🧬</div>
                      <p>CS7阶段细胞类型分布</p>
                      <p style={{ fontSize: "0.9rem", color: "var(--text-secondary)" }}>
                        早期发育阶段
                      </p>
                    </div>
                  )}
                </div>
              )}
            </div>

            {/* CS8阶段细胞类型分布 */}
            <div className="visualization-container">
              <h3 style={{ marginBottom: "1rem", textAlign: "center" }}>CS8 - 中期发育阶段</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>正在加载CS8数据...</p>
                    </div>
                  ) : (
                    <div>
                      <div style={{ fontSize: "2rem", marginBottom: "1rem" }}>🧬</div>
                      <p>CS8阶段细胞类型分布</p>
                      <p style={{ fontSize: "0.9rem", color: "var(--text-secondary)" }}>
                        中期发育阶段
                      </p>
                    </div>
                  )}
                </div>
              )}
            </div>

            {/* CS9阶段细胞类型分布 */}
            <div className="visualization-container">
              <h3 style={{ marginBottom: "1rem", textAlign: "center" }}>CS9 - 后期发育阶段</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>正在加载CS9数据...</p>
                    </div>
                  ) : (
                    <div>
                      <div style={{ fontSize: "2rem", marginBottom: "1rem" }}>🧬</div>
                      <p>CS9阶段细胞类型分布</p>
                      <p style={{ fontSize: "0.9rem", color: "var(--text-secondary)" }}>
                        后期发育阶段
                      </p>
                    </div>
                  )}
                </div>
              )}
            </div>
          </div>
        </div>

                {/* 第二部分和第三部分：并排显示 */}
        <div className="section">
          <h2>流动分析与基因时序</h2>
          <p className="flat-section-description">
            左侧桑基图展示细胞类型在发育阶段间的流动模式，右侧线图显示选定基因在各细胞类型中的占比变化。
          </p>
          
          <div
            style={{
              display: "grid",
              gridTemplateColumns: "1fr 1fr",
              gap: "2rem",
              marginBottom: "3rem"
            }}
          >
            {/* 桑基图部分 */}
            <div className="visualization-container">
              <h3 style={{ marginBottom: "1rem", textAlign: "center" }}>细胞类型演进流动图</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"
                    />
                    {/* 发育阶段标签行 */}
                    <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 - 早期
                      </div>
                      <div style={{ textAlign: "center", flex: 1 }}>
                        CS8 - 中期
                      </div>
                      <div style={{ textAlign: "right", flex: 1, paddingRight: "20px" }}>
                        CS9 - 后期
                      </div>
                    </div>
                  </>
                ) : (
                  <div className="no-data-message">
                    {sankeyLoading || loadingCS7 || loadingCS8 || loadingCS9 ? (
                      <div>
                        <div style={{ fontSize: "2rem", marginBottom: "1rem" }}>🔄</div>
                        <p>正在构建流动图...</p>
                      </div>
                    ) : (
                      <div>
                        <div style={{ fontSize: "2rem", marginBottom: "1rem" }}>📊</div>
                        <p>等待数据加载</p>
                      </div>
                    )}
                  </div>
                )}
              </div>
            </div>

            {/* 基因时序分析部分 */}
            <div className="visualization-container">
              <h3 style={{ marginBottom: "1rem", textAlign: "center" }}>基因时序分析</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="选择基因进行时序分析"
                  style={{ padding: "0.5rem", borderRadius: "0.3rem", border: "1px solid #ccc", fontSize: "0.9rem" }}
                >
                  <option value="">请选择基因</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>正在加载时序数据...</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: '发育阶段',
                      legendPosition: 'middle',
                      legendOffset: 40
                    }}
                    axisLeft={{
                      tickSize: 5,
                      tickPadding: 5,
                      tickRotation: 0,
                      legend: '占比 (%)',
                      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>请选择基因查看时序变化</p>
                  </div>
                )}
              </div>
            </div>
          </div>
        </div>

                 <div className="section">
           <h2>使用说明</h2>
           <div style={{ display: "grid", gridTemplateColumns: "repeat(auto-fit, minmax(300px, 1fr))", gap: "2rem", marginTop: "1.5rem" }}>
             <div>
               <h3>📅 时间序列对比</h3>
               <p>三个3D图分别展示CS7、CS8、CS9三个发育阶段的细胞类型分布，可以直观地观察细胞类型随时间的变化。</p>
             </div>
             <div>
               <h3>🌊 桑基图解读</h3>
               <p>桑基图显示细胞类型的演进流动：节点大小表示占比，连接线显示相同细胞类型在不同阶段间的延续性，颜色保持一致便于追踪。</p>
             </div>
             <div>
               <h3>📊 基因时序分析</h3>
               <p>选择特定基因后，堆叠面积图显示该基因在各细胞类型中的占比变化，帮助理解基因在发育过程中的表达模式演进。</p>
             </div>
             <div>
               <h3>🖱️ 交互操作</h3>
               <p>在3D视图中使用鼠标拖拽旋转视角，滚轮缩放，点击数据点查看详细信息。在桑基图和堆叠图中悬停查看具体数值。</p>
             </div>
             <div>
               <h3>🎨 颜色编码</h3>
               <p>相同颜色代表相同的细胞类型，在所有图表中保持一致，便于追踪特定细胞类型的变化和流动模式。</p>
             </div>
             <div>
               <h3>🔍 发育模式识别</h3>
               <p>通过对比三个阶段和不同基因，可以识别细胞类型的出现、消失、迁移和分化模式，以及基因表达的动态变化。</p>
             </div>
             <div>
               <h3>📈 百分比分析</h3>
               <p>桑基图和堆叠图中的百分比表示各细胞类型在对应发育阶段中的相对丰度，帮助理解细胞类型组成的动态平衡。</p>
             </div>
             <div>
               <h3>🧬 基因选择策略</h3>
               <p>选择不同的发育相关基因（如SOX2、NANOG等），观察其在胚胎发育过程中的表达变化模式和细胞类型特异性。</p>
             </div>
           </div>
         </div>
      </div>
    </div>
  );
};

export default TemporalAnalysis; 