// -*- coding: utf-8 -*-
// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.

using System;
using System.Linq;
using Microsoft.Quantum.Simulation.Core;
using Microsoft.Quantum.Simulation.Simulators;
using Microsoft.Quantum.Simulation.Simulators.QCTraceSimulators;
using Microsoft.Quantum.Chemistry;
using Microsoft.Extensions.Logging;
using System.Diagnostics;
using Microsoft.Quantum.Chemistry.OrbitalIntegrals;
using Microsoft.Quantum.Chemistry.Fermion;
using Microsoft.Quantum.Chemistry.QSharpFormat;
using McMaster.Extensions.CommandLineUtils;
using System.IO;
using Microsoft.Quantum.Chemistry.Broombridge;

// This loads a Hamiltonian from file and performs quantum phase estimation on
// - Jordan–Wigner Trotter step
// - Jordan–Wigner optimized Trotter step
// - Jordan–Wigner Qubitization iterate

namespace Microsoft.Quantum.Chemistry.Samples
{

    class Program
    {
        public static int Main(string[] args) =>
            CommandLineApplication.Execute<Program>(args);

        public enum DataFormat
        {
            LiQuiD, Broombridge
        }

        [Option(Description = "Format to use when loading data.")]
        public DataFormat Format { get; } = DataFormat.Broombridge;

        [Option(Description = "Path to data to be loaded.")]
        public string Path { get; } = System.IO.Path.Combine(
            "..", "IntegralData", "YAML", "lih_sto-3g_0.800_int.yaml"
        );

        [Option("--n-bits", Description = "Number of bits of precision to report from phase estimation.")]
        public int NBits { get; } = 10;

        [Option("--n-reps", Description = "Number of repetitions to use to find the minimum energy.")]
        public int NRepetitions { get; } = 5;

        [Option("--skip-jw", Description = "Skips simulating phase estimation with a Jordan–Wigner transformed Trotter–Suzuki approximation.")]
        public bool SkipJordanWigner { get; } = false;

        public bool RunJordanWigner => !SkipJordanWigner;

        [Option(Description = "Step size to use in the Trotter–Suzuki approximation.")]
        public double StepSize { get; } = 0.4;

        [Option("--skip-opt-jw", Description = "Skips simulating phase estimation with an optimized Jordan–Wigner transformed Trotter–Suzuki approximation.")]
        public bool SkipOptimizedJordanWigner { get; } = false;
        public bool RunOptimizedJordanWigner => !SkipOptimizedJordanWigner;

        [Option("--skip-qubitized-jw", Description = "Skips simulating phase estimation with a Jordan–Wigner transformed qubitization representation.")]
        public bool SkipQubitizedJordanWigner { get; } = false;
        public bool RunQubitizedJordanWigner => !SkipQubitizedJordanWigner;

        static void ParseLiQuiD(string path, out int nElectrons, out FermionHamiltonian hamiltonian)
        {
            using var reader = File.OpenText(path);
            nElectrons = 2;
            hamiltonian = LiQuiDSerializer
                          .Deserialize(reader)
                          .Single()
                          .OrbitalIntegralHamiltonian
                          .ToFermionHamiltonian(IndexConvention.UpDown);
        }

        static void ParseBroombridge(string path, out int nElectrons, out FermionHamiltonian hamiltonian)
        {
            using var reader = File.OpenText(path);
            var description = BroombridgeSerializer
                              .Deserialize(reader)
                              .Single();
            nElectrons = description.NElectrons;
            hamiltonian = description
                          .OrbitalIntegralHamiltonian
                          .ToFermionHamiltonian(IndexConvention.UpDown);
        }

        void OnExecute()
        {
            var logger = new LoggerFactory().CreateLogger<Program>();

            // Directory containing integral data generated by Microsoft.
            // Example Liquid data format files
            // spell-checker:words tzvp
            /*
            "h2_sto3g_4.dat" // 4 SO
            "B_sto6g.dat" // 10 SO
            "Be_sto6g_10.dat" // 10 SO
            "h2o_sto6g_14.dat" // 14 SO
            "h2s_sto6g_22.dat" // 22 SO
            "co2_sto3g_30.dat" // 30 SO
            "co2_p321_54.dat" // 54 SO
            "fe2s2_sto3g.dat" // 110 SO
            "nitrogenase_tzvp_54.dat" // 108 SO
            */
            


            // Read Hamiltonian terms from file.
            // Stopwatch for logging time to process file.
            Stopwatch stopWatch = new Stopwatch();


            #region For loading data in the format consumed by Liquid.
            stopWatch.Start();
            
            logger.LogInformation($"Processing {Path}...");
            int nElectrons;
            FermionHamiltonian fermionHamiltonian;
            if (Format == DataFormat.Broombridge)
            {
                ParseBroombridge(Path, out nElectrons, out fermionHamiltonian);
            }
            else
            {
                ParseLiQuiD(Path, out nElectrons, out fermionHamiltonian);
            }

            logger.LogInformation($"Load took {stopWatch.ElapsedMilliseconds}ms.");
            stopWatch.Restart();
            #endregion


            // Logs spin-orbital data in Logger.Message.
            logger.LogInformation(fermionHamiltonian.ToString());

            // Process Hamiltonian to obtain Jordan–Wigner representation.
            var jordanWignerEncoding = fermionHamiltonian.ToPauliHamiltonian(Paulis.QubitEncoding.JordanWigner);

            // Create input wavefunction.
            var wavefunction = fermionHamiltonian.CreateHartreeFockState(nElectrons);

            // Alternately, use wavefunction contained in problem description,
            // if available
            // var wavefunction = problemDescription.Wavefunctions["label"].ToIndexing(IndexConvention.UpDown);

            // Logs Jordan–Wigner representation data in Logger.Message.
            logger.LogInformation(jordanWignerEncoding.ToString());

            logger.LogInformation("End read file");

            // We begin by making an instance of the simulator that we will use to run our Q# code.
            using (var sim = new QuantumSimulator())
            {
                // Package hamiltonian and wavefunction data into a format
                // consumed by Q#.
                var qSharpData = QSharpFormat.Convert.ToQSharpFormat(
                    jordanWignerEncoding.ToQSharpFormat(),
                    wavefunction.ToQSharpFormat()
                );

                #region Calling into Q#

                if (RunJordanWigner)
                {
                    for (int i = 0; i < NRepetitions; i++)
                    {
                        var (phaseEst, energyEst) = TrotterEstimateEnergy.Run(sim, qSharpData, NBits, StepSize).Result;
                        logger.LogInformation($"Trotter simulation. phase: {phaseEst}; energy {energyEst}");
                    }
                }
                if (RunOptimizedJordanWigner)
                {
                    for (int i = 0; i < NRepetitions; i++)
                    {
                        var (phaseEst, energyEst) = OptimizedTrotterEstimateEnergy.Run(sim, qSharpData, NBits - 1, StepSize).Result;
                        logger.LogInformation($"Optimized Trotter simulation. phase {phaseEst}; energy {energyEst}");
                    }
                }
                if (RunQubitizedJordanWigner)
                {
                    for (int i = 0; i < NRepetitions; i++)
                    {
                        var (phaseEst, energyEst) = QubitizationEstimateEnergy.Run(sim, qSharpData, NBits - 2).Result;
                        logger.LogInformation($"Qubitization simulation. phase: {phaseEst}; energy {energyEst}");
                    }
                }

                #endregion
            }


            if (System.Diagnostics.Debugger.IsAttached)
            {
                System.Console.ReadLine();
            }

        }
    }
}

