import ai21
import os

ai21.api_key = os.getenv("AI21_API_KEY")


def text_completion(model, prompt, numResults, maxTokens, temperature, topKReturn, topP):
  response = ai21.Completion.execute(
    model=model,
    prompt=prompt,
    numResults=numResults,
    maxTokens=maxTokens,
    temperature=temperature,
    topKReturn=topKReturn,
    topP=topP,
    presencePenalty={
      "scale": 1,
      "applyToNumbers": True,
      "applyToPunctuations": True,
      "applyToStopwords": True,
      "applyToWhitespaces": True,
      "applyToEmojis": True
    },
    countPenalty={
      "scale": 1,
      "applyToNumbers": True,
      "applyToPunctuations": True,
      "applyToStopwords": True,
      "applyToWhitespaces": True,
      "applyToEmojis": True
    },
    frequencyPenalty={
      "scale": 1,
      "applyToNumbers": True,
      "applyToPunctuations": True,
      "applyToStopwords": True,
      "applyToWhitespaces": True,
      "applyToEmojis": True
    },
    stopSequences=[]
  )
  return response.suggestions[0].text

def chat(model, messages, numResults, maxTokens, temperature, topKReturn, topP):
  response = ai21.Chat.execute(
    model=model,
    messages=messages,
    numResults=numResults,
    maxTokens=maxTokens,
    temperature=temperature,
    topKReturn=topKReturn,
    topP=topP,
    presencePenalty={
      "scale": 1,
      "applyToNumbers": True,
      "applyToPunctuations": True,
      "applyToStopwords": True,
      "applyToWhitespaces": True,
      "applyToEmojis": True
    },
    countPenalty={
      "scale": 1,
      "applyToNumbers": True,
      "applyToPunctuations": True,
      "applyToStopwords": True,
      "applyToWhitespaces": True,
      "applyToEmojis": True
    },
    frequencyPenalty={
      "scale": 1,
      "applyToNumbers": True,
      "applyToPunctuations": True,
      "applyToStopwords": True,
      "applyToWhitespaces": True,
      "applyToEmojis": True
    },
    stopSequences=[]
  )
  return response.suggestions[0].text

def GEC(text):
  response =  ai21.GEC.execute(text=text)
  l = len(response.corrections)
  for i in range(l):
    sug = response.corrections[i].suggestion
    start = response.corrections[i].startIndex
    end = response.corrections[i].endIndex
    text = text.replace(
      text[start:end],
      sug
    )
  return text

def summarize(text):
  response = ai21.Summarize.execute(source=text, sourceType="TEXT")
  return response.summary

def improvements(text):
  response = ai21.Improvements.execute(
    text=text,
    types=[
      'fluency',
      'vocabulary/specificity',
      'vocabulary/variety',
      'clarity/short-sentences',
      'clarity/conciseness'
    ]
  )
  l = len(response.improvements)
  for i in range(l):
    sug = response.improvements[i].suggestions[0]
    start = response.improvements[i].startIndex
    end = response.improvements[i].endIndex
    text = text.replace(
      text[start:end],
      sug
    )

  return text

def paraphrase(text):
  response = ai21.Paraphrase.execute(text=text, style="general")
  return response.suggestions[0].text

def contextual_answer(context, question):
  response = ai21.Answer.execute(context=context, question=question)
  return response.suggestions[0].text