gendata

import org.apache.spark.sql.{DataFrame, SparkSession}
import org.apache.spark.sql.functions._
import org.apache.spark.sql.types._

/**
 * Klasa generująca dane testowe dla ShopEligibility
 */
class ShopEligibilityTestDataGenerator(spark: SparkSession) {
  import spark.implicits._

  // Stałe dla naszych testów
  val SHOP_GOOD = "SHOP001" // Sklep spełniający wszystkie kryteria
  val SHOP_BAD = "SHOP002"  // Sklep naruszający wszystkie reguły

  // Miesiące testowe (ostatnie 4 miesiące)
  val MONTHS = Array(202101, 202102, 202103, 202104)

  /**
   * Generowanie tabeli lrx_panel
   */
  def generatePanelData(): DataFrame = {
    val data = Seq(
      (SHOP_GOOD),
      (SHOP_BAD)
    )

    spark.createDataFrame(data.map(Tuple1.apply)).toDF("shop_id")
  }

  /**
   * Generowanie tabeli dim_pack
   */
  def generateDimPackData(): DataFrame = {
    val data = Seq(
      // Leki na receptę
      ("PACK001", "RX", 20210101, 20211231),
      ("PACK002", "RX", 20210101, 20211231),
      ("PACK003", "RX", 20210101, 20211231),
      // Leki bez recepty
      ("PACK004", "OTC", 20210101, 20211231),
      ("PACK005", "OTC", 20210101, 20211231)
    )

    spark.createDataFrame(data).toDF("pack_id", "basket_4_code", "day_id_from", "day_id_to")
  }

  /**
   * Generowanie tabeli lrx_sales
   */
  def generateSalesData(): DataFrame = {
    // Dla dobrego sklepu - dużo transakcji, dużo pacjentów, stabilne wartości RX
    val goodShopData = (for {
      month <- MONTHS
      // 30 pacjentów na miesiąc
      patientId <- 1 to 30
      // 2 transakcje na pacjenta
      trnOffset <- 0 to 1
      // 2 rodzaje paczek na transakcję (1 na receptę, 1 bez)
      packIds = if (trnOffset == 0) Seq("PACK001", "PACK004") else Seq("PACK002", "PACK005")
      packId <- packIds

      // Dzień w miesiącu
      day = month * 100 + 15

      // ID transakcji
      trnId = s"TRN_${month}_${patientId}_${trnOffset}"

      // ID przepisującego (tylko dla RX)
      prescriberId = if (packId.startsWith("PACK00") && packId.endsWith("1") || packId.endsWith("2") || packId.endsWith("3"))
                      Some(s"DOC_${patientId % 10 + 1}") else None
    } yield {
      (SHOP_GOOD, month, s"PAT_$patientId", trnId, packId, day, prescriberId.orNull)
    }).toSeq

    // Dla złego sklepu - mało transakcji, mało pacjentów, duże wahania RX
    val badShopData = (for {
      month <- MONTHS
      // Zmienna liczba pacjentów (tylko 10 w ostatnim miesiącu, ale rosnąca w czasie)
      patientMax = month match {
        case 202101 => 40  // Dużo w pierwszym miesiącu
        case 202102 => 5   // Gwałtowny spadek
        case 202103 => 30  // Gwałtowny wzrost
        case 202104 => 10  // Znowu spadek
      }
      patientId <- 1 to patientMax

      // Ograniczona liczba transakcji w ostatnim miesiącu, aby złamać regułę LOW_TRANSACTIONS_NUMBER
      if !(month == 202104 && patientId > 0)  // W ostatnim miesiącu będzie 0 transakcji

      // Pozostałe dane
      trnId = s"TRN_${month}_${patientId}_0"
      packId = if (patientId % 2 == 0) "PACK003" else "PACK005"
      day = month * 100 + 15
      prescriberId = if (packId == "PACK003") Some(s"DOC_${patientId % 5 + 1}") else None
    } yield {
      (SHOP_BAD, month, s"PAT_$patientId", trnId, packId, day, prescriberId.orNull)
    }).toSeq

    // Łączymy dane i tworzymy DataFrame
    val allData = goodShopData ++ badShopData

    spark.createDataFrame(allData).toDF(
      "shop_id", "month_id", "patient_id", "trn_id", "pack_id", "day_id", "prescriber_id"
    )
  }

  /**
   * Generuje wszystkie dane testowe i zapisuje je jako tabele tymczasowe
   */
  def generateAllTestData(): Unit = {
    val panel = generatePanelData()
    val dimPack = generateDimPackData()
    val sales = generateSalesData()

    // Tworzenie tabel tymczasowych
    panel.createOrReplaceTempView("lrx_panel")
    dimPack.createOrReplaceTempView("dim_pack")
    sales.createOrReplaceTempView("lrx_sales")

    println(s"Wygenerowano dane testowe:")
    println(s"  - Panel: ${panel.count()} sklepy")
    println(s"  - DimPack: ${dimPack.count()} paczek")
    println(s"  - Sales: ${sales.count()} rekordów sprzedaży")
  }
}