import { create } from 'zustand';
import { temporal } from 'zundo';

export interface ColumnMetadata {
  name: string;
  type: string;
  native_type: string;
}

export interface Outlier {
  index: number;
  reasons: string[];
}

export type ModelType = 'linear' | 'logistic' | 'polynomial' | 'exponential';

export interface AnalysisResults {
  r_squared: number;
  adj_r_squared: number | null;
  p_values: Record<string, number>;
  coefficients: Record<string, number>;
  std_errors: Record<string, number>;
  sample_size: number;
  residuals: number[];
  fit_plot: { real: number, pred: number }[]; // Changed to real vs pred structure
}

interface GridState {
  data: any[];
  columns: ColumnMetadata[];
  isLoading: boolean;
  error: string | null;
  modifiedCells: Set<string>;
  detectedOutliers: Record<string, Outlier[]>;
  multivariateOutliers: Outlier[];  // New: Separate storage for multivariate outliers
  excludedRows: Set<number>;
  activeInsightColumn: string | null;  // Column currently being viewed in InsightPanel

  // Analysis Config
  targetVariable: string | null;
  selectedFeatures: string[];
  modelType: ModelType;
  polyDegree: number;
  includeInteractions: boolean;
  analysisResults: AnalysisResults | null;

  setLoading: (loading: boolean) => void;
  setData: (data: any[], columns: ColumnMetadata[]) => void;
  setError: (error: string | null) => void;
  updateColumn: (oldName: string, newMetadata: Partial<ColumnMetadata>) => void;
  updateCellValue: (rowIndex: number, colName: string, value: any) => void;
  setOutliers: (colName: string, outliers: Outlier[]) => void;
  setMultivariateOutliers: (outliers: Outlier[]) => void;  // New: Setter for multivariate outliers
  toggleRowExclusion: (rowIndex: number) => void;
  excludeBatch: (rowIndices: number[]) => void;
  setActiveInsightColumn: (colName: string | null) => void;
  setTargetVariable: (name: string | null) => void;
  setSelectedFeatures: (names: string[]) => void;
  setModelType: (type: ModelType) => void;
  setPolyDegree: (degree: number) => void;
  setIncludeInteractions: (include: boolean) => void;
  setAnalysisResults: (results: AnalysisResults | null) => void;
  getCleanData: () => any[];
  reset: () => void;
}

export const useGridStore = create<GridState>()(
  temporal(
    (set, get) => ({
      data: [],
      columns: [],
      isLoading: false,
      error: null,
      modifiedCells: new Set(),
      detectedOutliers: {},
      multivariateOutliers: [],  // New: Initialize multivariate outliers
      excludedRows: new Set(),
      activeInsightColumn: null,  // Initialize active insight column
      targetVariable: null,
      selectedFeatures: [],
      modelType: 'linear',
      polyDegree: 2,
      includeInteractions: false,
      analysisResults: null,

      setLoading: (isLoading) => set({ isLoading }),
      setData: (data, columns) => set({
        data,
        columns,
        isLoading: false,
        error: null,
        modifiedCells: new Set(),
        detectedOutliers: {},
        multivariateOutliers: [],  // Reset multivariate outliers
        excludedRows: new Set(),
        activeInsightColumn: null,  // Reset active insight column
        targetVariable: null,
        selectedFeatures: [],
        modelType: 'linear',
        polyDegree: 2,
        includeInteractions: false,
        analysisResults: null
      }),
      setError: (error) => set({ error, isLoading: false }),

      updateColumn: (oldName, newMetadata) => set((state) => {
        const newColumns = state.columns.map((col) => col.name === oldName ? { ...col, ...newMetadata } : col);
        let newData = state.data;
        if (newMetadata.name && newMetadata.name !== oldName) {
          newData = state.data.map((row) => {
            const newRow = { ...row };
            newRow[newMetadata.name!] = newRow[oldName];
            delete newRow[oldName];
            return newRow;
          });
        }
        return { columns: newColumns, data: newData };
      }),

      updateCellValue: (rowIndex, colName, value) => set((state) => {
        const newData = [...state.data];
        newData[rowIndex] = { ...newData[rowIndex], [colName]: value };
        const newModified = new Set(state.modifiedCells);
        newModified.add(`${rowIndex}:${colName}`);
        return { data: newData, modifiedCells: newModified };
      }),

      setOutliers: (colName, outliers) => set((state) => ({
        detectedOutliers: { ...state.detectedOutliers, [colName]: outliers }
      })),

      setMultivariateOutliers: (outliers) => set({ multivariateOutliers: outliers }),  // New: Setter for multivariate outliers

      setActiveInsightColumn: (colName) => set({ activeInsightColumn: colName }),  // Set which column's insights are being viewed

      toggleRowExclusion: (rowIndex) => set((state) => {
        const newExcluded = new Set(state.excludedRows);
        const isRemoving = newExcluded.has(rowIndex);

        if (isRemoving) {
          newExcluded.delete(rowIndex);
        } else {
          newExcluded.add(rowIndex);
        }

        // Si on exclut une ligne, retirer ses outliers des listes
        let newDetectedOutliers = state.detectedOutliers;
        let newMultivariateOutliers = state.multivariateOutliers;

        if (!isRemoving) {
          // Retirer des detectedOutliers
          newDetectedOutliers = { ...state.detectedOutliers };
          Object.keys(newDetectedOutliers).forEach(colName => {
            newDetectedOutliers[colName] = newDetectedOutliers[colName].filter(
              outlier => outlier.index !== rowIndex
            );
          });

          // Retirer des multivariateOutliers
          newMultivariateOutliers = state.multivariateOutliers.filter(
            outlier => outlier.index !== rowIndex
          );
        }

        return {
          excludedRows: newExcluded,
          detectedOutliers: newDetectedOutliers,
          multivariateOutliers: newMultivariateOutliers
        };
      }),

      excludeBatch: (rowIndices) => set((state) => {
        const newExcluded = new Set(state.excludedRows);
        rowIndices.forEach(idx => newExcluded.add(idx));

        // Retirer les outliers exclus de detectedOutliers
        const newDetectedOutliers = { ...state.detectedOutliers };
        Object.keys(newDetectedOutliers).forEach(colName => {
          newDetectedOutliers[colName] = newDetectedOutliers[colName].filter(
            outlier => !newExcluded.has(outlier.index)
          );
        });

        // Retirer les outliers exclus de multivariateOutliers
        const newMultivariateOutliers = state.multivariateOutliers.filter(
          outlier => !newExcluded.has(outlier.index)
        );

        return {
          excludedRows: newExcluded,
          detectedOutliers: newDetectedOutliers,
          multivariateOutliers: newMultivariateOutliers
        };
      }),

      setTargetVariable: (targetVariable) => set({ targetVariable }),
      setSelectedFeatures: (selectedFeatures) => set({ selectedFeatures }),
      setModelType: (modelType) => set({ modelType }),
      setPolyDegree: (polyDegree) => set({ polyDegree }),
      setIncludeInteractions: (includeInteractions) => set({ includeInteractions }),
      setAnalysisResults: (analysisResults) => set({ analysisResults }),

      getCleanData: () => {
        const { data, excludedRows } = get();
        return data.filter((_, index) => !excludedRows.has(index));
      },

      reset: () => set({
        data: [], columns: [], isLoading: false, error: null,
        modifiedCells: new Set(), detectedOutliers: {}, multivariateOutliers: [], excludedRows: new Set(),
        targetVariable: null, selectedFeatures: [], modelType: 'linear',
        polyDegree: 2, includeInteractions: false,
        analysisResults: null
      }),
    }),
    {
      partialize: (state) => {
        const {
          data,
          columns,
          modifiedCells,
          excludedRows,
          detectedOutliers,
          multivariateOutliers,
          activeInsightColumn,
          targetVariable,
          selectedFeatures,
          modelType,
          polyDegree,
          includeInteractions
        } = state;
        return {
          data,
          columns,
          modifiedCells,
          excludedRows,
          detectedOutliers,
          multivariateOutliers,
          activeInsightColumn,
          targetVariable,
          selectedFeatures,
          modelType,
          polyDegree,
          includeInteractions
        };
      },
    }
  )
);
