{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "8c9c7f3e-accc-4fdd-8d91-ae603f5d21c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "5564717a-681b-4ee4-a852-a6c637013a67",
   "metadata": {},
   "outputs": [],
   "source": [
    "def preprocess_dataset(file_path):\n",
    "    df = pd.read_csv(file_path, parse_dates=[0])\n",
    "\n",
    "    result = {}\n",
    "    for ticker in df['tck_iem_cd'].unique():\n",
    "        ticker_data = df[df['tck_iem_cd'] == ticker]\n",
    "        result[ticker] = ticker_data[['gts_iem_ong_pr', 'gts_iem_end_pr', 'gts_iem_hi_pr', 'gts_iem_low_pr']].values.tolist()\n",
    "\n",
    "    return result\n",
    "\n",
    "# def preprocess_dataset(file_path):\n",
    "#     df = pd.read_csv(file_path, parse_dates=[0])\n",
    "\n",
    "#     tickers_of_interest = ['NVDA', 'TSLA', 'AAPL']\n",
    "#     result = {}\n",
    "\n",
    "#     for ticker in tickers_of_interest:\n",
    "#         if ticker in df['tck_iem_cd'].unique():\n",
    "#             ticker_data = df[df['tck_iem_cd'] == ticker]\n",
    "#             result[ticker] = ticker_data[['gts_iem_ong_pr', 'gts_iem_end_pr', 'gts_iem_hi_pr', 'gts_iem_low_pr']].values.tolist()\n",
    "\n",
    "#     return result\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "0a52ffb2-e28c-49d3-afac-87805ac00a78",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Usage\n",
    "file_path = \"NASDAQ_DT_FC_STK_QUT.csv\"\n",
    "preprocessed_data = preprocess_dataset(file_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f50a8f9b-578d-4538-aed0-51eb11d76af6",
   "metadata": {},
   "source": [
    "# DataLoader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "ef2475e0-44dc-4ca1-b574-3912aafcbb90",
   "metadata": {},
   "outputs": [],
   "source": [
    "def encode_data(data):\n",
    "    encoded_data = {}\n",
    "    for ticker, sequences in data.items():\n",
    "        encoded_sequences = []\n",
    "        for sequence in sequences:\n",
    "            encoded_sequence = []\n",
    "            for day in sequence:\n",
    "                open_price, close_price, high_price, low_price = day\n",
    "                body_size = close_price - open_price\n",
    "                upper_shadow = high_price - max(open_price, close_price)\n",
    "                lower_shadow = min(open_price, close_price) - low_price\n",
    "                encoded_sequence.append([body_size, upper_shadow, lower_shadow])\n",
    "            encoded_sequences.append(encoded_sequence)\n",
    "        encoded_data[ticker] = encoded_sequences\n",
    "    return encoded_data\n",
    "\n",
    "\n",
    "\n",
    "# 노말라이즈 한 값을 기준으로 change 계산한거 같음-> 에러\n",
    "\n",
    "def create_labels(data):\n",
    "    labels = {}\n",
    "    for ticker, sequences in data.items():\n",
    "        sequence_labels = []\n",
    "        for i, sequence in enumerate(sequences[:-1]): # Excluding the last sequence since we don't have the next day data for it\n",
    "            current_end_price = sequence[-1][1]\n",
    "            next_end_price = sequences[i+1][-1][1]\n",
    "            price_change = (next_end_price - current_end_price) / current_end_price\n",
    "            if price_change > 0.02:\n",
    "                label = 0\n",
    "            elif price_change < -0.02:\n",
    "                label = 1\n",
    "            else:\n",
    "                label = 2\n",
    "            sequence_labels.append(label)\n",
    "        labels[ticker] = sequence_labels\n",
    "    return labels\n",
    "\n",
    "\n",
    "import torch.nn as nn\n",
    "\n",
    "class TransformerPatternRecognition(nn.Module):\n",
    "    def __init__(self, d_model, nhead, num_encoder_layers, num_classes):\n",
    "        super(TransformerPatternRecognition, self).__init__()\n",
    "        self.transformer = nn.Transformer(d_model, nhead, num_encoder_layers)\n",
    "        self.fc = nn.Linear(d_model, num_classes)\n",
    "\n",
    "    def forward(self, x):\n",
    "        output = self.transformer(x, x)  # Using x as both source and target\n",
    "        output = self.fc(output[:, -1, :])  # We only need the last day's output for classification\n",
    "        return output\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "aa98ef82-3199-485f-aa57-612b22130a88",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_98718/2083546567.py:28: RuntimeWarning: divide by zero encountered in scalar divide\n",
      "  price_change = (next_end_price - current_end_price) / current_end_price\n",
      "/tmp/ipykernel_98718/2083546567.py:28: RuntimeWarning: invalid value encountered in scalar divide\n",
      "  price_change = (next_end_price - current_end_price) / current_end_price\n"
     ]
    }
   ],
   "source": [
    "def chunk_sequences(data, window_size=5):\n",
    "    \"\"\"\n",
    "    Chunks the data into sequences of a given window size.\n",
    "    \"\"\"\n",
    "    chunks = []\n",
    "    for i in range(len(data) - window_size + 1):\n",
    "        chunks.append(data[i:i+window_size])\n",
    "    return chunks\n",
    "\n",
    "def normalize_chunk(chunk):\n",
    "    \"\"\"\n",
    "    Normalizes a given chunk using the first day's start price.\n",
    "    The first day's start price is set to 0, while other prices are log-scaled.\n",
    "    \"\"\"\n",
    "    first_day_start_price = chunk[0][0]\n",
    "    normalized_chunk = []\n",
    "\n",
    "    for day in chunk:\n",
    "        normalized_day = [\n",
    "            np.log(day[i] / first_day_start_price) if i == 0 and day_idx == 0 else \n",
    "            (np.log(day[i] / first_day_start_price) if day[i] != 0 else -np.inf)\n",
    "            for day_idx, i in enumerate(range(4))\n",
    "        ]\n",
    "        normalized_chunk.append(normalized_day)\n",
    "\n",
    "    # Setting first day's start price to 0\n",
    "    normalized_chunk[0][0] = 0\n",
    "    \n",
    "    return normalized_chunk\n",
    "\n",
    "def preprocess_data_for_training(data, window_size=5):\n",
    "    # Existing functions: chunk_sequences, normalize_chunk, encode_data, create_labels\n",
    "\n",
    "    training_data = []\n",
    "\n",
    "    for ticker, prices in data.items():\n",
    "        chunks = chunk_sequences(prices, window_size)\n",
    "        normalized_chunks = [normalize_chunk(chunk) for chunk in chunks]\n",
    "        encoded_chunks = encode_data({'temp': normalized_chunks})['temp']  # Using a temp key to fit the function's format\n",
    "        \n",
    "        # Combine normalized and encoded data\n",
    "        combined_chunks = [\n",
    "            [day + encoded_day for day, encoded_day in zip(norm_chunk, encoded_chunk)] \n",
    "            for norm_chunk, encoded_chunk in zip(normalized_chunks, encoded_chunks)\n",
    "        ]\n",
    "\n",
    "        # Create labels for these chunks\n",
    "        labels_for_chunks = create_labels({'temp': normalized_chunks})['temp']\n",
    "\n",
    "        # Add to the training data\n",
    "        for combined_chunk, label in zip(combined_chunks, labels_for_chunks):\n",
    "            training_data.append((combined_chunk, label))\n",
    "\n",
    "    return training_data\n",
    "\n",
    "\n",
    "training_data = preprocess_data_for_training(preprocessed_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "ada18c84-a17f-47db-ace8-5bb8d62083dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import DataLoader, TensorDataset\n",
    "\n",
    "# Create DataLoader\n",
    "def get_dataloader(training_data, batch_size=32):\n",
    "    input_tensors = torch.tensor([item[0] for item in training_data], dtype=torch.float32)\n",
    "    label_tensors = torch.tensor([item[1] for item in training_data], dtype=torch.long)\n",
    "    dataset = TensorDataset(input_tensors, label_tensors)\n",
    "    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)\n",
    "    return dataloader\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "14c09114-a4aa-42f6-8cec-c176ae5c1fbd",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/1000, Loss: nan, Accuracy: 38.10%\n",
      "Epoch 2/1000, Loss: nan, Accuracy: 34.78%\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "KeyboardInterrupt\n",
      "\n"
     ]
    }
   ],
   "source": [
    "import torch.optim as optim\n",
    "import torch\n",
    "\n",
    "# Parameters\n",
    "d_model = 7 # 4 normalized values (start, end, high, low) + 3 encoded values (body_size, upper_shadow, lower_shadow)\n",
    "nhead = 1\n",
    "num_encoder_layers = 2\n",
    "num_classes = 3\n",
    "learning_rate = 0.001\n",
    "num_epochs = 100\n",
    "\n",
    "# Model, Loss, Optimizer\n",
    "model = TransformerPatternRecognition(d_model, nhead, num_encoder_layers, num_classes)\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = optim.Adam(model.parameters(), lr=learning_rate)\n",
    "\n",
    "def calculate_accuracy(outputs, labels):\n",
    "    # Extract predicted class (the one with the highest probability)\n",
    "    _, predicted = torch.max(outputs, 1)\n",
    "    \n",
    "    # Filter out class 2 from both predicted and actual labels\n",
    "    valid_indices = labels != 2\n",
    "    valid_predictions = predicted[valid_indices]\n",
    "    valid_labels = labels[valid_indices]\n",
    "    \n",
    "    correct = (valid_predictions == valid_labels).sum().item()\n",
    "    total = valid_labels.size(0)\n",
    "    \n",
    "    return 100 * correct / total if total != 0 else 0\n",
    "\n",
    "# Training loop with batching\n",
    "dataloader = get_dataloader(training_data)\n",
    "for epoch in range(num_epochs):\n",
    "    total_loss = 0.0\n",
    "    for inputs, labels in dataloader:\n",
    "        optimizer.zero_grad()\n",
    "        outputs = model(inputs)\n",
    "        loss = criterion(outputs, labels)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        \n",
    "        total_loss += loss.item()\n",
    "    \n",
    "    avg_loss = total_loss / len(dataloader)\n",
    "    accuracy = calculate_accuracy(outputs, labels)\n",
    "    print(f\"Epoch {epoch+1}/{num_epochs}, Loss: {avg_loss:.4f}, Accuracy: {accuracy:.2f}%\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0c3e2e31-844f-498d-9f55-3763c990e73b",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b3c1371-4b40-460b-9b60-419266fcec38",
   "metadata": {},
   "outputs": [],
   "source": [
    "# threshold를 높여서 해당 추론한것에 대해서만 정확도 계산 -> 정확도 식 재작성"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "745f25ee-0331-49b5-b793-758746638d5b",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "prac",
   "language": "python",
   "name": "prac"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}