import * as tf  from '@tensorflow/tfjs'
import { f7 } from 'framework7-vue'

let model = null

export default {
  methods: {
    async loadModel(weights, preload) {
      if (model && model.modelURL == weights) {
        return model
      } else if (model) {
        tf.dispose(model)
      }
      model = await tf.loadGraphModel(weights)
      const [modelWidth, modelHeight] = model.inputs[0].shape.slice(1, 3)
      /*****************
       * If preloading then run model
       * once on fake data to preload
       * weights for a faster response
       *****************/
      if (preload) {
        const dummyT = tf.ones([1,modelWidth,modelHeight,3])
        model.predict(dummyT)
      }
      return model
    },
    async localDetect(imageData) {
      console.time('pre-process')
      const [modelWidth, modelHeight] = model.inputs[0].shape.slice(1, 3)
      let gTense = null
      const input = tf.tidy(() => {
        gTense = tf.image.rgbToGrayscale(tf.image.resizeBilinear(tf.browser.fromPixels(imageData), [modelWidth, modelHeight])).div(255.0).expandDims(0)
        return tf.concat([gTense,gTense,gTense],3)
      })
      tf.dispose(gTense)
      console.timeEnd('pre-process')

      console.time('run prediction')
      const res = model.predict(input)
      const tRes = tf.transpose(res,[0,2,1])
      const rawRes = tRes.arraySync()[0]
      console.timeEnd('run prediction')

      console.time('post-process')
      const outputSize = res.shape[1]
      let rawBoxes = []
      let rawScores = []

      for (var i = 0; i < rawRes.length; i++) {
        var getScores = rawRes[i].slice(4)
        if (getScores.every( s => s < .05)) { continue }
        var getBox = rawRes[i].slice(0,4)
        var boxCalc = [
          (getBox[0] - (getBox[2] / 2)) / modelWidth,
          (getBox[1] - (getBox[3] / 2)) / modelHeight,
          (getBox[0] + (getBox[2] / 2)) / modelWidth,
          (getBox[1] + (getBox[3] / 2)) / modelHeight,
        ]
        rawBoxes.push(boxCalc)
        rawScores.push(getScores)
      }

      if (rawBoxes.length > 0) {
        const tBoxes = tf.tensor2d(rawBoxes)
        let tScores = null
        let resBoxes = null
        let validBoxes = []
        let structureScores = null
        let boxes_data = []
        let scores_data = []
        let classes_data = []
        for (var c = 0; c < outputSize - 4; c++) {
          structureScores = rawScores.map(x => x[c])
          tScores = tf.tensor1d(structureScores)
          resBoxes = await tf.image.nonMaxSuppressionAsync(tBoxes,tScores,10,0.5,.05)
          validBoxes = resBoxes.dataSync()
          tf.dispose(resBoxes)
          if (validBoxes) {
            boxes_data.push(...rawBoxes.filter( (_, idx) => validBoxes.includes(idx)))
            var outputScores = structureScores.filter( (_, idx) => validBoxes.includes(idx))
            scores_data.push(...outputScores)
            classes_data.push(...outputScores.fill(c))
          }
        }

        validBoxes = []
        tf.dispose(tBoxes)
        tf.dispose(tScores)
        tf.dispose(tRes)
        const valid_detections_data = classes_data.length
        var output = {
          detections: []
        }
        for (var i =0; i < valid_detections_data; i++) {
          var [dLeft, dTop, dRight, dBottom] = boxes_data[i]
          output.detections.push({
            "top": dTop,
            "left": dLeft,
            "bottom": dBottom,
            "right": dRight,
            "label": this.detectorLabels[classes_data[i]].name,
            "confidence": scores_data[i] * 100
          })
        }
      }
      tf.dispose(res)
      tf.dispose(input)
      console.timeEnd('post-process')

      return output || { detections: [] }
    },
    getRemoteLabels() {
      var self = this
      var modelURL = `http://${this.serverSettings.address}:${this.serverSettings.port}/detectors`
      var xhr = new XMLHttpRequest()
      xhr.open("GET", modelURL)
      xhr.setRequestHeader('Content-Type', 'application/json')
      xhr.timeout = 10000
      xhr.ontimeout = this.remoteTimeout
      xhr.onload = function () {
        if (this.status !== 200) {
          console.log(xhr.response)
          const errorResponse = JSON.parse(xhr.response)
          f7.dialog.alert(`ALVINN has encountered an error: ${errorResponse.error}`)
          return
        }
        var detectors = JSON.parse(xhr.response).detectors
        var findLabel = detectors
          .find( d => { return d.name == self.detectorName } )?.labels
          .filter( l => { return l != "" } ).sort()
          .map( l => { return {'name': l, 'detect': true} } )
        self.detectorLabels = findLabel || []
      }
      xhr.onerror = function (e) {
        f7.dialog.alert('ALVINN has encountered an unknown server error')
        return
      }
      
      xhr.send()
    },
    remoteDetect() {
      var self = this
      var modelURL = `http://${this.serverSettings.address}:${this.serverSettings.port}/detect`
      var xhr = new XMLHttpRequest()
      xhr.open("POST", modelURL)
      xhr.timeout = 10000
      xhr.ontimeout = this.remoteTimeout
      xhr.setRequestHeader('Content-Type', 'application/json')
      xhr.onload = function () {
        self.detecting = false
        if (this.status !== 200) {
          console.log(xhr.response)
          const errorResponse = JSON.parse(xhr.response)
          f7.dialog.alert(`ALVINN has encountered an error: ${errorResponse.error}`)
          return;
        }
        self.resultData = JSON.parse(xhr.response)
        self.uploadDirty = true
      }
      
      var doodsData = {
        "detector_name": this.detectorName,
        "detect": {
          "*": 1
        },
        "data": this.imageView.src.split(',')[1]
      }

      xhr.send(JSON.stringify(doodsData))
    },
    remoteTimeout () {
      this.detecting = false
      f7.dialog.alert('No connection to remote ALVINN instance.  Please check app settings.')
    },
    async videoFrameDetect (vidData) {
      await this.loadModel(this.miniLocation)
      const [modelWidth, modelHeight] = model.inputs[0].shape.slice(1, 3)
      const imCanvas = this.$refs.image_cvs
      const imageCtx = imCanvas.getContext("2d")
      const target = this.$refs.target_image
      await tf.nextFrame();
      imCanvas.width = imCanvas.clientWidth
      imCanvas.height = imCanvas.clientHeight
      imageCtx.clearRect(0,0,imCanvas.width,imCanvas.height)
      var imgWidth
      var imgHeight
      const imgAspect = vidData.width / vidData.height
      const rendAspect = imCanvas.width / imCanvas.height
      if (imgAspect >= rendAspect) {
        imgWidth = imCanvas.width
        imgHeight = imCanvas.width / imgAspect
      } else {
        imgWidth = imCanvas.height * imgAspect
        imgHeight = imCanvas.height
      }
      while (this.videoAvailable) {
        console.time('frame-process')
        try {
          const input = tf.tidy(() => {
            return tf.image.resizeBilinear(tf.browser.fromPixels(vidData), [modelWidth, modelHeight]).div(255.0).expandDims(0)
          })
          const res = model.predict(input)
          const rawRes = tf.transpose(res,[0,2,1]).arraySync()[0]

          let rawCoords = []
          if (rawRes) {
            for (var i = 0; i < rawRes.length; i++) {
              let getScores = rawRes[i].slice(4) 
              if (getScores.some( s => s > .5)) {
                let foundTarget = rawRes[i].slice(0,2)
                foundTarget.push(Math.max(...getScores))
                rawCoords.push(foundTarget)
              }
            }

            imageCtx.clearRect(0,0,imCanvas.width,imCanvas.height)
            for (var coord of rawCoords) {
              console.log(`x: ${coord[0]}, y: ${coord[1]}`)
              let pointX = (imCanvas.width - imgWidth) / 2 + (coord[0] / modelWidth) * imgWidth -5
              let pointY = (imCanvas.height - imgHeight) / 2 + (coord[1] / modelHeight) * imgHeight -5
              imageCtx.globalAlpha = coord[2]
              imageCtx.drawImage(target, pointX, pointY, 20, 20)
            }
          }
        } catch (e) {
          console.log(e)
        }
        console.timeEnd('frame-process')
        await tf.nextFrame();
      }
    }
  }
}