import { h, Component } from 'preact'
import { bindActionCreators } from 'redux'
import { connect } from 'react-redux'
import util from '../../../util'

import * as pix2pixhdActions from '../pix2pixhd.actions'
import * as pix2pixhdTasks from '../pix2pixhd.tasks'

import {
  Loading,
  FileList, FileRow,
  Select, SelectGroup, Group, Param, Button,
  TextInput, NumberInput,
  CurrentTask, TaskList,
  AugmentationGrid,
} from '../../../common'
import DatasetForm from '../../../dataset/dataset.form'
import NewDatasetForm from '../../../dataset/dataset.new'
import UploadStatus from '../../../dataset/upload.status'

import DatasetComponent from '../../../dataset/dataset.component'
import SequenceEditor from './sequence.editor'

import pix2pixhdModule from '../pix2pixhd.module'

class Pix2PixHDTrain extends Component {
  constructor(props){
    super(props)
    this.handleChange = this.handleChange.bind(this)
    this.state = {
      checkpoint_name: 'PLACEHOLDER',
      epoch: 'latest',
      augment_name: '',
      augment_take: 100,
      augment_make: 20,
    }
  }
  componentDidMount(){
    const id = this.props.match.params.id || localStorage.getItem('pix2pixhd.last_id')
    const dataset = localStorage.getItem('pix2pixhd.last_dataset')
    console.log('load dataset:', id, dataset)
    const { match, pix2pixhd, actions } = this.props
    if (id === 'new') return
    if (id) {
      if (parseInt(id)) localStorage.setItem('pix2pixhd.last_id', id)
      if (! pix2pixhd.folder || pix2pixhd.folder.id !== id) {
        actions.load_directories(id)
      }
    } else {
      this.props.history.push('/pix2pixhd/new/')
    }
    if (dataset) {
      this.setState({ checkpoint_name: dataset })
    }
  }
  componentDidUpdate(prevProps, prevState){
    if (prevState.checkpoint_name !== this.state.checkpoint_name) {
      localStorage.setItem('pix2pixhd.last_dataset', this.state.checkpoint_name)
      this.setState({ epoch: 'latest' })
      this.props.actions.list_epochs(this.state.checkpoint_name)
      this.props.actions.count_dataset(this.state.checkpoint_name)
    }
  }
  handleChange(name, value){
    console.log('name', name, 'value', value)
    this.setState({ [name]: value })
  }
  interrupt(){
    this.props.actions.queue.stop_task('gpu')
  }
  render(){
    if (this.props.pix2pixhd.loading) {
      return <Loading progress={this.props.pix2pixhd.progress} />
    }
    const { pix2pixhd, match, history, queue } = this.props
    const { folderLookup, datasetLookup } = (pix2pixhd.data || {})
    const folder = (folderLookup || {})[pix2pixhd.folder_id] || {}
    // console.log(pix2pixhd)

    const checkpointGroups = Object.keys(folderLookup).map(id => {
      const folder = this.props.pix2pixhd.data.folderLookup[id]
      if (folder.name === 'results') return
      const datasets = folder.datasets.map(name => {
        const dataset = datasetLookup[name]
        if (dataset.checkpoints.length) {
          return name
        }
        return null
      }).filter(n => !!n)
      return {
        name: folder.name,
        options: datasets.sort(),
      }
    }).filter(n => !!n && !!n.options.length).sort((a,b) => a.name.localeCompare(b.name))

    // console.log('state', this.props.pix2pixhd.data.epochs)
    // console.log(this.state.checkpoint_name, this.state.epoch)
    // console.log(queue)
    return (
      <div className='app pix2pixhd'>
        <div className='heading'>
          <h1>pix2pixhd training</h1>
        </div>
        <div className='columns'>
          <div className='column'>
            <Group title='Dataset'>
              <SelectGroup
                name='checkpoint_name'
                title='Dataset'
                options={checkpointGroups}
                onChange={this.handleChange}
                placeholder='Pick a dataset'
                value={this.state.checkpoint_name}
              />
              <Select
                title="Epoch"
                name="epoch"
                options={this.props.pix2pixhd.data.epochs}
                onChange={this.handleChange}
                value={this.state.epoch}
              />
            </Group>
          </div>
        </div>
        <div>
          <Group title='Sequence Editor'>
            <SequenceEditor
              module={pix2pixhdModule}
              checkpoint={this.props.pix2pixhd.checkpoint}
            />
          </Group>
        </div>
        <div className='columns'>
          <div className='column'>
            <Group title='Augmentation Grid'>
              <AugmentationGrid
                checkpoint={this.props.pix2pixhd.checkpoint}
                take={[1,2,3,4,5,10,15,20,25,50,75,100,200,300,400,500,1000]}
                make={[1,2,3,4,5,10,15,20,25,50,75,100,200,]}
                onAugment={(augment_take, augment_make) => {
                  this.props.remote.augment_task(this.state.checkpoint_name, {
                    ...this.state,
                    augment_take,
                    augment_make,
                  })
                }}
                onTrain={() => {
                  this.props.remote.train_task(this.state.checkpoint_name, pix2pixhd.folder_id, 1)
                  setTimeout(() => { // auto-generate epoch demo
                    this.props.remote.augment_task(this.state.checkpoint_name, {
                      ...this.state,
                      augment_take: 10,
                      augment_make: 149,
                      no_symlinks: true,
                      mov: true,
                      folder_id: this.props.pix2pixhd.data.resultsFolder.id
                    })
                  }, 250)
                }}
              />
            </Group>
            <Group title='Augment'>
              <NumberInput
                name="augment_take"
                title="Pick N random frames"
                value={this.state.augment_take}
                onChange={this.handleChange}
                type="int"
                min="1"
                max="1000"
              />
              <NumberInput
                name="augment_make"
                title="Generate N recursively"
                value={this.state.augment_make}
                onChange={this.handleChange}
                type="int"
                min="1"
                max="1000"
              />
              <Button
                title="Augment dataset"
                value="Augment"
                onClick={() => this.props.remote.augment_task(this.state.checkpoint_name, this.state)}
              />
              <Button
                title="Make a movie without augmenting"
                value="Generate"
                onClick={() => {
                  this.props.remote.augment_task(this.state.checkpoint_name, {
                    ...this.state,
                    no_symlinks: true,
                    mov: true,
                    folder_id: this.props.pix2pixhd.data.resultsFolder.id
                  })
                }}
              />
            </Group>

            <Group title='Status'>
              <Button
                title="GPU"
                value={this.props.runner.gpu.status === 'IDLE' ? "Idle" : "Interrupt"}
                onClick={() => this.interrupt()}
              />
              <CurrentTask />
            </Group>
          </div>
          <div className='column'>
            <Group title='Upcoming Tasks'>
              <TaskList tasks={queue.queue.map(id => queue.tasks[id])} sort="date asc" />
            </Group>
          </div>
        </div>
      </div>
    )
  }
}

const mapStateToProps = state => ({
  pix2pixhd: state.module.pix2pixhd,
  runner: state.system.runner,
  queue: state.queue,
})

const mapDispatchToProps = (dispatch, ownProps) => ({
  actions: bindActionCreators(pix2pixhdActions, dispatch),
  remote: bindActionCreators(pix2pixhdTasks, dispatch),
})

export default connect(mapStateToProps, mapDispatchToProps)(Pix2PixHDTrain)