// IMPORT LIBRARIES TOOLS
import { pipeline, env } from 'https://cdn.jsdelivr.net/npm/@xenova/transformers@2.10.1';

// skip local model check
env.allowLocalModels = false;

// GLOBAL VARIABLES

var PREPROMPT = `Please continue the story, and fill any [MASK] with your own words:`
// let PREPROMPT = `Please complete the phrase and fill in any [MASK]: `
var PROMPT_INPUT = `The [BLANK] has a job as a [MASK] but...` // a field for writing or changing a text value
var promptField // an html element to hold the prompt
var outText // an html element to hold the results
var blanksArray = [] // an empty list to store all the variables we enter to modify the prompt
// e.g. ["woman", "man", "non-binary person"]

// RUN TEXT-GEN MODEL

async function textGenTask(pre, prompt, blanks){
  console.log('text-gen task initiated')
  
  // Create concatenated prompt array including preprompt and all variable prompts
  let promptArray = []
  promptArray.push(pre) // add preprompt to the list of prompts

  // Fill in blanks from our sample prompt and make new prompts using our variable list 'blanksArray'
  blanks.forEach(b => {
    let p = prompt.replace('[BLANK]', b) // replace the string segment with an item from the blanksArray
    promptArray.push(p) // add the new prompt to the list we created
  })

  // create combined fill prompt
  let INPUT = promptArray.toString()
  console.log(INPUT)
  // let INPUT = pre + prompt // simple concatenated input
  // let INPUT = prompt // basic prompt input

  // PICK MODEL 
  let MODEL = 'Xenova/flan-alpaca-large'

  // MODELS LIST
  // - Xenova/bloom-560m
  // - Xenova/distilgpt2
  // - Xenova/LaMini-Cerebras-256M
  // - Xenova/gpt-neo-125M // not working well
  // - Xenova/llama2.c-stories15M // only fairytails
  // - webml/TinyLlama-1.1B-Chat-v1.0
  // - Xenova/TinyLlama-1.1B-Chat-v1.0
  // - Xenova/flan-alpaca-large //text2text

  
  // const pipe = await pipeline('text-generation', MODEL) //different task type, also for text generation
  const pipe = await pipeline('text2text-generation', MODEL)

  var hyperparameters = { max_new_tokens: 200, top_k: 90, repetition_penalty: 1.5 }
    // setting hyperparameters
    // max_new_tokens: 256, top_k: 50, temperature: 0.7, do_sample: true, no_repeat_ngram_size: 2, num_return_sequences: 2 (must be 1?)

  // change model run to iterative for each prompt generated locally — will be more expensive??
  // promptArray.forEach(async i => {} //this was a loop to wrap model run multiple times
  
  // RUN INPUT THROUGH MODEL, 
  var out = await pipe(INPUT, hyperparameters)

  console.log(await out)
  console.log('text-gen task completed')
  
  // PARSE RESULTS as a list of outputs, two different ways depending on the model
  
  // parsing of output
  // await out.forEach(o => {
  //   console.log(o)
  //   OUTPUT_LIST.push(o.generated_text)
  // })
  
  // alternate format for parsing, for chat model type
  // await out.choices.forEach(o => {
  //   console.log(o)
  //   OUTPUT_LIST.push(o.message.content)
  // })

  let OUTPUT_LIST = out[0].generated_text //not a list anymore just one result
  // OUTPUT_LIST.push(out[0].generated_text)

  console.log(OUTPUT_LIST)
  console.log('text-gen parsing complete')

  return await OUTPUT_LIST
  // return await out
}

// RUN FILL-IN MODEL
async function fillInTask(input){
  console.log('fill-in task initiated')

  // MODELS LIST
  // - Xenova/bert-base-uncased

  const pipe = await pipeline('fill-mask', 'Xenova/bert-base-uncased');
  
  var out = await pipe(input);

  console.log(await out) // yields { score, sequence, token, token_str } for each result

  let OUTPUT_LIST = [] // a blank array to store the results from the model

  // parsing of output
  await out.forEach(o => {
    console.log(o) // yields { score, sequence, token, token_str } for each result
    OUTPUT_LIST.push(o.sequence) // put only the full sequence in a list
  })
  
  console.log(await OUTPUT_LIST)
  console.log('fill-in task completed')

  // return await out
  return await OUTPUT_LIST
}

//// p5.js Instance

new p5(function (p5){
  p5.setup = function(){
      p5.noCanvas()
      console.log('p5 instance loaded')
      makeTextDisplay()
      makeFields()
      makeButtons()
    }

  p5.draw = function(){
      // 
  }

  function makeTextDisplay(){
    const title = p5.createElement('h1','p5.js Critical AI Prompt Battle')
    const intro = p5.createP(`This tool lets you explore several AI prompts results at once.`) 
    p5.createP(`Use it to explore what models 'know' about various concepts, communities, and cultures. For more information on prompt programming and critical AI, see [Tutorial & extra info][TO-DO][XXX]`)
  }

  function makeFields(){
    promptField = p5.createInput(PROMPT_INPUT) // turns the string into an input; now access the text via PROMPT_INPUT.value()
    promptField.size(700)
    promptField.attribute('label', `Write a text prompt with one [MASK] that the model will fill in.`)
    p5.createP(promptField.attribute('label'))
    promptField.addClass("prompt")

    const fieldsDiv = p5.createDiv()
    fieldsDiv.id('fieldsDiv')

    addField()
  }

  function makeButtons(){
    // press to run model
    const submitButton = p5.createButton("SUBMIT")
    submitButton.size(170)
    submitButton.class('submit')
    submitButton.mousePressed(displayResults)

    // press to add more blanks to fill in
    const addButton = p5.createButton("more blanks")
    addButton.size(170)
    // addButton.position(220,500)
    addButton.mousePressed(addField)
    
    // also make results placeholder
    const outHeader = p5.createElement('h3',"Results")
    outText = p5.createP('').id('results')
  }

  async function displayResults(){
    console.log('submitButton pressed')

    // insert waiting dots into results space of interface
    outText.html('...', false)

    // GRAB CURRENT FIELD INPUTS FROM PROMPT & BLANKS
    PROMPT_INPUT = promptField.value() // grab update to the prompt if it's been changed
    console.log("latest prompt: ", PROMPT_INPUT)

    let blanksValues = []
    blanksArray.forEach(b => {
      blanksValues.push(b.value())
    })
    console.log(blanksValues)
    // let blanksValues = blanksArray.map(b => b.value)    
    
    // call the function that runs the model for the task of your choice here 
    // make sure to use the PROMPT_INPUT as a parameter, or also the PREPROMPT if valid for that task
    let outs = await textGenTask(PREPROMPT, PROMPT_INPUT, blanksValues)
    console.log(outs)

    // insert the model outputs into the paragraph
    await outText.html(outs, false) // false replaces text instead of appends
  }

  function addField(){
      let f = p5.createInput("")
      f.class("blank")
      f.parent("#fieldsDiv")

      blanksArray.push(f)
      console.log("made variable field")
      
      // Cap the number of fields, avoids token limit in prompt
      let blanks = document.querySelectorAll(".blank")
      if (blanks.length > 5){
          console.log(blanks.length)
          addButton.style('visibility','hidden')
      }
  }
  
});