# This code is part of quantastica.qiskit_toaster
#
# (C) Copyright Quantastica 2019. 
# https://quantastica.com/
#
# This code is licensed under the Apache License, Version 2.0. You may
# obtain a copy of this license in the LICENSE.txt file in the root directory
# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
#
# Any modifications or derivative works of this code must retain this
# copyright notice, and modified files need to carry a notice indicating
# that they have been altered from the originals.

from concurrent import futures
import logging
import sys
import subprocess
import json

from quantastica.qconvert import qobj_to_toaster

from qiskit.providers import BaseJob, JobStatus, JobError
from qiskit.qobj import validate_qobj_against_schema
from qiskit.result import Result

logger = logging.getLogger(__name__)


class ToasterJob(BaseJob):
    _MINQTOASTERVERSION = '0.9.9'
    
    _executor = futures.ThreadPoolExecutor()

    def __init__(self, backend, job_id, qobj, toasterpath, 
                 getstates = False):
        super().__init__(backend, job_id)
        self._result = None
        self._qobj = qobj
        self._future = None
        self._toasterpath = toasterpath;
        self._getstates = getstates

    def submit(self):
        if self._future is not None:
            raise JobError("We have already submitted the job!")

        logger.debug("submit...")
        validate_qobj_against_schema(self._qobj)
        self._future = self._executor.submit(self._run_with_qtoaster)

    def wait(self, timeout=None):
        if self.status() is JobStatus.RUNNING :
            futures.wait([self._future], timeout);
            if self._future.exception() :
                raise self._future.exception()

    def result(self, timeout=None):
        self.wait(timeout)
        return Result.from_dict(self._result);

    def cancel(self):
        return
        # return self._future.cancel()

    def status(self):

        if self._future is None:
            _status = JobStatus.INITIALIZING
        elif self._future.running():
            _status = JobStatus.RUNNING
        elif self._future.cancelled():
            _status = JobStatus.CANCELLED
        elif self._future.done():
            _status = JobStatus.DONE if self._future.exception() is None else JobStatus.ERROR
        else:
            _status = JobStatus.INITIALIZING

        return _status

    def backend(self):
        """Return the instance of the backend used for this job."""
        return self._backend
    
    @staticmethod
    def _qtoaster_version_to_int(versionstring):
        parts = versionstring.split(".")
        verint = 0
        if len(parts) >= 3 :
            verint = int(parts[0])*1000000 + int(parts[1])*1000 + int(parts[2])
        return verint

    @classmethod
    def _check_qtoaster_version(cls, versionstring):
        verint = cls._qtoaster_version_to_int(versionstring)
        minverint = cls._qtoaster_version_to_int(cls._MINQTOASTERVERSION)
        if verint < minverint :
            return False
        return True

    @staticmethod
    def _convert_counts(counts):
        ret = dict()
        for key in counts:
            nicekey = key.replace(' ','')
            nicekey = hex(int(nicekey,2))
            ret[nicekey] = counts[key];
        return ret

    def _run_with_qtoaster(self):
        qobj_dict = self._qobj.to_dict()
        if self._getstates :
            shots = 1
        else:
            shots = qobj_dict['config']['shots']

        converted = qobj_to_toaster(qobj_dict, { "all_experiments": False })

        args = [
            self._toasterpath,
            "-",
            "-r",
            "counts",
            "-s",
            "%d"%shots
        ]
        if self._getstates:
            args.append("-r")
            args.append("state")

        logger.info("Running q-toaster with following params:")
        logger.info(args)
        proc = subprocess.run(
            args, 
            input=converted.encode(),
            stdout=subprocess.PIPE )

        qtoasterjson = proc.stdout

        resultraw = json.loads(qtoasterjson)
        logger.debug("Raw results from toaster:\r\n%s",resultraw)
        if isinstance(resultraw , dict) :
            rawversion = resultraw.get('qtoaster_version')
        else:
            rawversion = "0.0.0"

        if ToasterJob._check_qtoaster_version(rawversion) == False :
            raise ValueError(
                "Unsupported qtoaster_version, got '%s' - minimum expected is '%s'.\n\rPlease update your q-toaster to latest version"%(rawversion,self._MINQTOASTERVERSION))

        counts = self._convert_counts(resultraw['counts'])
        qobjid = qobj_dict['qobj_id']
        qobj_header = qobj_dict['header']
        exp_dict = qobj_dict['experiments'][0]
        exp_header = exp_dict['header']
        expname = exp_header['name']
        qubit_count = exp_dict['config']['n_qubits']
        statevector = resultraw.get('statevector')
        data = dict()
        data['counts'] = counts;
        if statevector is not None and len(statevector) > 0:
            data['statevector'] = statevector

        self._result = {
            'success': True, 
            'backend_name': qobj_header['backend_name'], 
            'qobj_id': qobjid ,
            'backend_version': rawversion, 
            'header': qobj_header,
            'job_id': self._job_id, 
            'results': [
                {
                    'success': True, 
                    'meas_level': 2, 
                    'shots': shots, 
                    'data': data, 
                    'header': exp_header, 
                    'status': 'DONE', 
                    'time_taken': resultraw['time_taken'], 
                    'name': expname, 
                    'seed_simulator': 0
                }
                ], 
            'status': 'COMPLETED'
        }
