import type { AlgorithmDefinition, HeapNodeState, HeapState } from '@/lib/visualizations/types';

/**
 * Kth Largest Element in an Array (LeetCode 215)
 *
 * Input: nums = [3, 2, 1, 5, 6, 4], k = 2
 * Output: 5 (2nd largest element)
 *
 * Using a min-heap of size k:
 * - Push each element to heap
 * - If heap size > k, pop the smallest
 * - Root of min-heap = kth largest
 */

// Helper to create input array state
function createInputArray(
  values: number[],
  currentIndex?: number,
  processedIndices: number[] = []
) {
  return {
    id: 'input',
    label: 'Input Array',
    elements: values.map((value, index) => ({
      value,
      index,
      state: (index === currentIndex
        ? 'highlighted'
        : processedIndices.includes(index)
        ? 'dimmed'
        : 'normal') as 'normal' | 'highlighted' | 'dimmed' | 'success' | 'comparing',
    })),
  };
}

// Helper to create heap state
function createHeapState(
  elements: { value: number; state: HeapNodeState['state'] }[],
  label?: string
): HeapState {
  return {
    id: 'heap',
    label: label ?? 'Min-Heap (k=2)',
    heapType: 'min',
    maxSize: 2,
    elements: elements.map((el, index) => ({
      id: `heap-${index}`,
      value: el.value,
      index,
      state: el.state,
    })),
  };
}

export const kthLargestAlgorithm: AlgorithmDefinition = {
  id: 'kth-largest',
  title: 'Kth Largest Element (Min-Heap)',
  slug: 'kth-largest',
  pattern: {
    name: 'Heap',
    description:
      'Use a min-heap of size k to efficiently find the kth largest element in O(n log k) time.',
  },
  problemStatement:
    'Given an integer array nums and an integer k, return the kth largest element in the array. Use a heap-based approach instead of sorting.',
  intuition:
    'A min-heap of size k always keeps the k largest elements seen so far. The root (minimum) of this heap is exactly the kth largest overall. As we process elements, we push to the heap and pop if it exceeds size k.',
  code: {
    language: 'python',
    code: `import heapq

def find_kth_largest(nums: list[int], k: int) -> int:
    heap = []

    for num in nums:
        heapq.heappush(heap, num)  # Add element
        if len(heap) > k:
            heapq.heappop(heap)    # Remove smallest

    return heap[0]  # Root = kth largest`,
  },
  initialExample: {
    input: { nums: [3, 2, 1, 5, 6, 4], k: 2 },
    expected: 5,
  },
  steps: [
    // ==========================================
    // Phase 1: Problem (2 steps)
    // ==========================================
    {
      id: 'problem-1',
      phase: 'problem',
      explanation:
        'Find the 2nd largest element in [3, 2, 1, 5, 6, 4]. Sorting gives [1, 2, 3, 4, 5, 6], so 2nd largest is 5. But can we do better than O(n log n)?',
      dataState: {
        arrays: [createInputArray([3, 2, 1, 5, 6, 4])],
        pointers: [],
        variables: [{ id: 'k', name: 'k', value: 2 }],
        calculations: [],
        heaps: [createHeapState([])],
      },
    },
    {
      id: 'problem-2',
      phase: 'problem',
      explanation:
        'Brute force: Sort the array and return nums[n-k]. Time: O(n log n). With a heap, we can achieve O(n log k), which is faster when k << n.',
      dataState: {
        arrays: [createInputArray([3, 2, 1, 5, 6, 4])],
        pointers: [],
        variables: [
          { id: 'k', name: 'k', value: 2 },
          { id: 'n', name: 'n', value: 6 },
        ],
        calculations: [],
        heaps: [createHeapState([])],
      },
    },

    // ==========================================
    // Phase 2: Intuition (3 steps)
    // ==========================================
    {
      id: 'intuition-1',
      phase: 'intuition',
      explanation:
        'Key insight: If we maintain a collection of the k largest elements seen so far, the smallest among them is the kth largest overall.',
      dataState: {
        arrays: [createInputArray([3, 2, 1, 5, 6, 4])],
        pointers: [],
        variables: [{ id: 'k', name: 'k', value: 2 }],
        calculations: [],
        heaps: [createHeapState([])],
      },
    },
    {
      id: 'intuition-2',
      phase: 'intuition',
      explanation:
        'A min-heap is perfect: it gives O(log k) access to the minimum element. We keep exactly k elements; the root is always the kth largest.',
      dataState: {
        arrays: [createInputArray([3, 2, 1, 5, 6, 4])],
        pointers: [],
        variables: [{ id: 'k', name: 'k', value: 2 }],
        calculations: [],
        heaps: [createHeapState([
          { value: 5, state: 'highlighted' },
          { value: 6, state: 'normal' },
        ], 'Goal: min-heap with k=2 largest')],
      },
    },
    {
      id: 'intuition-3',
      phase: 'intuition',
      explanation:
        'Process: For each element, push to heap. If heap size exceeds k, pop the minimum. This evicts smaller elements, keeping only the k largest.',
      dataState: {
        arrays: [createInputArray([3, 2, 1, 5, 6, 4])],
        pointers: [],
        variables: [{ id: 'k', name: 'k', value: 2 }],
        calculations: [],
        heaps: [createHeapState([
          { value: 5, state: 'highlighted' },
          { value: 6, state: 'normal' },
        ])],
      },
    },

    // ==========================================
    // Phase 3: Pattern (2 steps)
    // ==========================================
    {
      id: 'pattern-1',
      phase: 'pattern',
      explanation:
        'Heap property: In a min-heap, parent ≤ children. Array index i has children at 2i+1 and 2i+2. Root at index 0 is always the minimum.',
      dataState: {
        arrays: [createInputArray([3, 2, 1, 5, 6, 4])],
        pointers: [],
        variables: [],
        calculations: [],
        heaps: [createHeapState([])],
      },
    },
    {
      id: 'pattern-2',
      phase: 'pattern',
      explanation:
        'Top-K pattern: Use min-heap of size k for "k largest", max-heap of size k for "k smallest". The root gives the kth element after processing all items.',
      dataState: {
        arrays: [createInputArray([3, 2, 1, 5, 6, 4])],
        pointers: [],
        variables: [{ id: 'k', name: 'k', value: 2 }],
        calculations: [],
        heaps: [createHeapState([])],
      },
    },

    // ==========================================
    // Phase 4: Code (3 steps)
    // ==========================================
    {
      id: 'code-1',
      phase: 'code',
      explanation:
        'Initialize an empty heap. We\'ll use Python\'s heapq which implements a min-heap.',
      codeLine: 4,
      codeHighlightLines: [4],
      dataState: {
        arrays: [createInputArray([3, 2, 1, 5, 6, 4])],
        pointers: [],
        variables: [{ id: 'k', name: 'k', value: 2 }],
        calculations: [],
        heaps: [createHeapState([])],
      },
    },
    {
      id: 'code-2',
      phase: 'code',
      explanation:
        'For each number: push to heap (heappush maintains heap property by bubbling up). If size > k: pop minimum (heappop bubbles down to restore property).',
      codeLine: 7,
      codeHighlightLines: [6, 7, 8, 9],
      dataState: {
        arrays: [createInputArray([3, 2, 1, 5, 6, 4])],
        pointers: [],
        variables: [{ id: 'k', name: 'k', value: 2 }],
        calculations: [],
        heaps: [createHeapState([])],
      },
    },
    {
      id: 'code-3',
      phase: 'code',
      explanation:
        'After processing all elements, the heap contains exactly k largest. Return heap[0] - the root, which is the minimum of the k largest = kth largest overall.',
      codeLine: 11,
      codeHighlightLines: [11],
      dataState: {
        arrays: [createInputArray([3, 2, 1, 5, 6, 4])],
        pointers: [],
        variables: [{ id: 'k', name: 'k', value: 2 }],
        calculations: [],
        heaps: [createHeapState([
          { value: 5, state: 'highlighted' },
          { value: 6, state: 'normal' },
        ])],
      },
    },

    // ==========================================
    // Phase 5: Execution (~25 steps)
    // ==========================================
    // Process 3
    {
      id: 'exec-1',
      phase: 'execution',
      explanation:
        'Start with empty heap. Process first element: 3.',
      codeLine: 6,
      dataState: {
        arrays: [createInputArray([3, 2, 1, 5, 6, 4], 0)],
        pointers: [],
        variables: [
          { id: 'k', name: 'k', value: 2 },
          { id: 'num', name: 'num', value: 3 },
        ],
        calculations: [],
        heaps: [createHeapState([])],
      },
    },
    {
      id: 'exec-2',
      phase: 'execution',
      explanation:
        'Push 3 to heap. Heap: [3]. Size = 1 ≤ k = 2, so no pop needed.',
      codeLine: 7,
      dataState: {
        arrays: [createInputArray([3, 2, 1, 5, 6, 4], 0)],
        pointers: [],
        variables: [
          { id: 'k', name: 'k', value: 2 },
          { id: 'size', name: 'size', value: 1 },
        ],
        calculations: [],
        heaps: [createHeapState([
          { value: 3, state: 'settled' },
        ])],
      },
    },

    // Process 2
    {
      id: 'exec-3',
      phase: 'execution',
      explanation:
        'Process next element: 2.',
      codeLine: 6,
      dataState: {
        arrays: [createInputArray([3, 2, 1, 5, 6, 4], 1, [0])],
        pointers: [],
        variables: [
          { id: 'k', name: 'k', value: 2 },
          { id: 'num', name: 'num', value: 2, previousValue: 3 },
        ],
        calculations: [],
        heaps: [createHeapState([
          { value: 3, state: 'normal' },
        ])],
      },
    },
    {
      id: 'exec-4',
      phase: 'execution',
      explanation:
        'Push 2 to heap. 2 < 3, so 2 bubbles up to become the new root. Heap: [2, 3].',
      codeLine: 7,
      dataState: {
        arrays: [createInputArray([3, 2, 1, 5, 6, 4], 1, [0])],
        pointers: [],
        variables: [
          { id: 'k', name: 'k', value: 2 },
          { id: 'size', name: 'size', value: 2 },
        ],
        calculations: [],
        heaps: [createHeapState([
          { value: 2, state: 'comparing' },
          { value: 3, state: 'swapping' },
        ])],
      },
    },
    {
      id: 'exec-5',
      phase: 'execution',
      explanation:
        'Heap size = 2 = k. No pop needed. Heap: [2, 3]. Root 2 would be kth largest so far.',
      codeLine: 8,
      dataState: {
        arrays: [createInputArray([3, 2, 1, 5, 6, 4], 1, [0])],
        pointers: [],
        variables: [
          { id: 'k', name: 'k', value: 2 },
          { id: 'size', name: 'size', value: 2 },
        ],
        calculations: [],
        heaps: [createHeapState([
          { value: 2, state: 'settled' },
          { value: 3, state: 'settled' },
        ])],
      },
    },

    // Process 1
    {
      id: 'exec-6',
      phase: 'execution',
      explanation:
        'Process next element: 1.',
      codeLine: 6,
      dataState: {
        arrays: [createInputArray([3, 2, 1, 5, 6, 4], 2, [0, 1])],
        pointers: [],
        variables: [
          { id: 'k', name: 'k', value: 2 },
          { id: 'num', name: 'num', value: 1, previousValue: 2 },
        ],
        calculations: [],
        heaps: [createHeapState([
          { value: 2, state: 'normal' },
          { value: 3, state: 'normal' },
        ])],
      },
    },
    {
      id: 'exec-7',
      phase: 'execution',
      explanation:
        'Push 1 to heap. 1 < 2, so 1 bubbles up. Heap becomes: [1, 3, 2].',
      codeLine: 7,
      dataState: {
        arrays: [createInputArray([3, 2, 1, 5, 6, 4], 2, [0, 1])],
        pointers: [],
        variables: [
          { id: 'k', name: 'k', value: 2 },
          { id: 'size', name: 'size', value: 3 },
        ],
        calculations: [],
        heaps: [createHeapState([
          { value: 1, state: 'comparing' },
          { value: 3, state: 'normal' },
          { value: 2, state: 'swapping' },
        ])],
      },
    },
    {
      id: 'exec-8',
      phase: 'execution',
      explanation:
        'Heap size = 3 > k = 2. Pop the minimum (1). After pop and bubble-down: [2, 3].',
      codeLine: 9,
      dataState: {
        arrays: [createInputArray([3, 2, 1, 5, 6, 4], 2, [0, 1])],
        pointers: [],
        variables: [
          { id: 'k', name: 'k', value: 2 },
          { id: 'size', name: 'size', value: 2, previousValue: 3 },
          { id: 'popped', name: 'popped', value: 1 },
        ],
        calculations: [],
        heaps: [createHeapState([
          { value: 1, state: 'removing' },
          { value: 3, state: 'normal' },
          { value: 2, state: 'normal' },
        ])],
      },
    },
    {
      id: 'exec-9',
      phase: 'execution',
      explanation:
        'After removing 1, heap restored to [2, 3]. Element 1 was too small to be in top-2.',
      codeLine: 9,
      dataState: {
        arrays: [createInputArray([3, 2, 1, 5, 6, 4], 2, [0, 1])],
        pointers: [],
        variables: [
          { id: 'k', name: 'k', value: 2 },
          { id: 'size', name: 'size', value: 2 },
        ],
        calculations: [],
        heaps: [createHeapState([
          { value: 2, state: 'settled' },
          { value: 3, state: 'settled' },
        ])],
      },
    },

    // Process 5
    {
      id: 'exec-10',
      phase: 'execution',
      explanation:
        'Process next element: 5.',
      codeLine: 6,
      dataState: {
        arrays: [createInputArray([3, 2, 1, 5, 6, 4], 3, [0, 1, 2])],
        pointers: [],
        variables: [
          { id: 'k', name: 'k', value: 2 },
          { id: 'num', name: 'num', value: 5, previousValue: 1 },
        ],
        calculations: [],
        heaps: [createHeapState([
          { value: 2, state: 'normal' },
          { value: 3, state: 'normal' },
        ])],
      },
    },
    {
      id: 'exec-11',
      phase: 'execution',
      explanation:
        'Push 5 to heap. 5 > 2 (parent), so no bubble-up needed. Heap: [2, 3, 5].',
      codeLine: 7,
      dataState: {
        arrays: [createInputArray([3, 2, 1, 5, 6, 4], 3, [0, 1, 2])],
        pointers: [],
        variables: [
          { id: 'k', name: 'k', value: 2 },
          { id: 'size', name: 'size', value: 3 },
        ],
        calculations: [],
        heaps: [createHeapState([
          { value: 2, state: 'normal' },
          { value: 3, state: 'normal' },
          { value: 5, state: 'comparing' },
        ])],
      },
    },
    {
      id: 'exec-12',
      phase: 'execution',
      explanation:
        'Heap size = 3 > k = 2. Pop minimum (2). Move 5 to root, bubble down: 5 > 3, so swap. Result: [3, 5].',
      codeLine: 9,
      dataState: {
        arrays: [createInputArray([3, 2, 1, 5, 6, 4], 3, [0, 1, 2])],
        pointers: [],
        variables: [
          { id: 'k', name: 'k', value: 2 },
          { id: 'popped', name: 'popped', value: 2 },
        ],
        calculations: [],
        heaps: [createHeapState([
          { value: 2, state: 'removing' },
          { value: 3, state: 'swapping' },
          { value: 5, state: 'swapping' },
        ])],
      },
    },
    {
      id: 'exec-13',
      phase: 'execution',
      explanation:
        'After pop: heap = [3, 5]. Now tracking top-2: {3, 5}. Root 3 = current kth largest candidate.',
      codeLine: 9,
      dataState: {
        arrays: [createInputArray([3, 2, 1, 5, 6, 4], 3, [0, 1, 2])],
        pointers: [],
        variables: [
          { id: 'k', name: 'k', value: 2 },
          { id: 'size', name: 'size', value: 2 },
        ],
        calculations: [],
        heaps: [createHeapState([
          { value: 3, state: 'settled' },
          { value: 5, state: 'settled' },
        ])],
      },
    },

    // Process 6
    {
      id: 'exec-14',
      phase: 'execution',
      explanation:
        'Process next element: 6.',
      codeLine: 6,
      dataState: {
        arrays: [createInputArray([3, 2, 1, 5, 6, 4], 4, [0, 1, 2, 3])],
        pointers: [],
        variables: [
          { id: 'k', name: 'k', value: 2 },
          { id: 'num', name: 'num', value: 6, previousValue: 5 },
        ],
        calculations: [],
        heaps: [createHeapState([
          { value: 3, state: 'normal' },
          { value: 5, state: 'normal' },
        ])],
      },
    },
    {
      id: 'exec-15',
      phase: 'execution',
      explanation:
        'Push 6 to heap. 6 > 3 (parent), no bubble-up. Heap: [3, 5, 6].',
      codeLine: 7,
      dataState: {
        arrays: [createInputArray([3, 2, 1, 5, 6, 4], 4, [0, 1, 2, 3])],
        pointers: [],
        variables: [
          { id: 'k', name: 'k', value: 2 },
          { id: 'size', name: 'size', value: 3 },
        ],
        calculations: [],
        heaps: [createHeapState([
          { value: 3, state: 'normal' },
          { value: 5, state: 'normal' },
          { value: 6, state: 'comparing' },
        ])],
      },
    },
    {
      id: 'exec-16',
      phase: 'execution',
      explanation:
        'Heap size = 3 > k. Pop minimum (3). After bubble-down: [5, 6]. Now tracking {5, 6}.',
      codeLine: 9,
      dataState: {
        arrays: [createInputArray([3, 2, 1, 5, 6, 4], 4, [0, 1, 2, 3])],
        pointers: [],
        variables: [
          { id: 'k', name: 'k', value: 2 },
          { id: 'popped', name: 'popped', value: 3 },
        ],
        calculations: [],
        heaps: [createHeapState([
          { value: 3, state: 'removing' },
          { value: 5, state: 'swapping' },
          { value: 6, state: 'normal' },
        ])],
      },
    },
    {
      id: 'exec-17',
      phase: 'execution',
      explanation:
        'Heap now [5, 6]. Root 5 = 2nd largest so far. The two largest seen are 6 and 5.',
      codeLine: 9,
      dataState: {
        arrays: [createInputArray([3, 2, 1, 5, 6, 4], 4, [0, 1, 2, 3])],
        pointers: [],
        variables: [
          { id: 'k', name: 'k', value: 2 },
          { id: 'size', name: 'size', value: 2 },
        ],
        calculations: [],
        heaps: [createHeapState([
          { value: 5, state: 'settled' },
          { value: 6, state: 'settled' },
        ])],
      },
    },

    // Process 4
    {
      id: 'exec-18',
      phase: 'execution',
      explanation:
        'Process last element: 4.',
      codeLine: 6,
      dataState: {
        arrays: [createInputArray([3, 2, 1, 5, 6, 4], 5, [0, 1, 2, 3, 4])],
        pointers: [],
        variables: [
          { id: 'k', name: 'k', value: 2 },
          { id: 'num', name: 'num', value: 4, previousValue: 6 },
        ],
        calculations: [],
        heaps: [createHeapState([
          { value: 5, state: 'normal' },
          { value: 6, state: 'normal' },
        ])],
      },
    },
    {
      id: 'exec-19',
      phase: 'execution',
      explanation:
        'Push 4 to heap. 4 < 5 (parent at index 0), so 4 bubbles up. Heap becomes [4, 6, 5].',
      codeLine: 7,
      dataState: {
        arrays: [createInputArray([3, 2, 1, 5, 6, 4], 5, [0, 1, 2, 3, 4])],
        pointers: [],
        variables: [
          { id: 'k', name: 'k', value: 2 },
          { id: 'size', name: 'size', value: 3 },
        ],
        calculations: [],
        heaps: [createHeapState([
          { value: 4, state: 'comparing' },
          { value: 6, state: 'normal' },
          { value: 5, state: 'swapping' },
        ])],
      },
    },
    {
      id: 'exec-20',
      phase: 'execution',
      explanation:
        'Heap size = 3 > k. Pop minimum (4). After pop: [5, 6]. 4 < 5, so it\'s not in top-2.',
      codeLine: 9,
      dataState: {
        arrays: [createInputArray([3, 2, 1, 5, 6, 4], 5, [0, 1, 2, 3, 4])],
        pointers: [],
        variables: [
          { id: 'k', name: 'k', value: 2 },
          { id: 'popped', name: 'popped', value: 4 },
        ],
        calculations: [],
        heaps: [createHeapState([
          { value: 4, state: 'removing' },
          { value: 6, state: 'swapping' },
          { value: 5, state: 'normal' },
        ])],
      },
    },
    {
      id: 'exec-21',
      phase: 'execution',
      explanation:
        'Final heap: [5, 6]. These are the 2 largest elements from the array.',
      codeLine: 9,
      dataState: {
        arrays: [createInputArray([3, 2, 1, 5, 6, 4], undefined, [0, 1, 2, 3, 4, 5])],
        pointers: [],
        variables: [
          { id: 'k', name: 'k', value: 2 },
          { id: 'size', name: 'size', value: 2 },
        ],
        calculations: [],
        heaps: [createHeapState([
          { value: 5, state: 'settled' },
          { value: 6, state: 'settled' },
        ])],
      },
    },

    // Return result
    {
      id: 'exec-22',
      phase: 'execution',
      explanation:
        'Return heap[0] = 5. The root of the min-heap (smallest of the k largest) is the kth largest element.',
      codeLine: 11,
      dataState: {
        arrays: [createInputArray([3, 2, 1, 5, 6, 4], undefined, [0, 1, 2, 3, 4, 5])],
        pointers: [],
        variables: [
          { id: 'k', name: 'k', value: 2 },
          { id: 'result', name: 'result', value: 5 },
        ],
        calculations: [],
        heaps: [createHeapState([
          { value: 5, state: 'highlighted' },
          { value: 6, state: 'normal' },
        ])],
      },
    },
    {
      id: 'exec-23',
      phase: 'execution',
      explanation:
        'Done! The 2nd largest element is 5. Time: O(n log k), Space: O(k). Much faster than sorting when k is small.',
      codeLine: 11,
      decision: {
        question: 'Why min-heap for kth largest?',
        answer: 'Min-heap lets us quickly remove the smallest among k elements',
        action: 'Root is always the kth largest after processing all elements',
      },
      dataState: {
        arrays: [createInputArray([3, 2, 1, 5, 6, 4], undefined, [0, 1, 2, 3, 4, 5])],
        pointers: [],
        variables: [
          { id: 'result', name: 'result', value: 5 },
          { id: 'time', name: 'Time', value: 'O(n log k)' },
          { id: 'space', name: 'Space', value: 'O(k)' },
        ],
        calculations: [],
        heaps: [createHeapState([
          { value: 5, state: 'highlighted' },
          { value: 6, state: 'normal' },
        ])],
      },
    },
  ],
};
