
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.cm as cm
import os
from .backend import execute, StatevectorSimulator, StabilizerSimulator
from .quantum import QuantumCircuit
try:
    import seaborn as sns
    _HAS_SNS = True
except ImportError:
    sns = None
    _HAS_SNS = False
try:
    from qutip import Qobj, ket2dm
    _HAS_QUTIP = True
except ImportError:
    Qobj = ket2dm = None
    _HAS_QUTIP = False
try:
    import imageio
    _HAS_IMAGEIO = True
except ImportError:
    imageio = None
    _HAS_IMAGEIO = False
try:
    import plotly.graph_objects as go
    from plotly.subplots import make_subplots
    _HAS_PLOTLY = True
except ImportError:
    go = None
    make_subplots = None
    _HAS_PLOTLY = False

def partial_trace(rho, keep):
    n = int(np.log2(rho.shape[0]))
    rho = rho.reshape([2]*(2*n))
    keep = set(keep)
    traced = []
    for i in range(n):
        if i not in keep:
            ri = i - len([j for j in traced if j < i])
            ci = (n+i) - len([j for j in traced if j < n+i])
            rho = np.trace(rho, axis1=ri, axis2=ci)
            traced += [i, n+i]
    return rho

def get_bloch_vector(rho):
    px = np.array([[0,1],[1,0]])
    py = np.array([[0,-1j],[1j,0]])
    pz = np.array([[1,0],[0,-1]])
    return (float(np.real(np.trace(rho@px))),
            float(np.real(np.trace(rho@py))),
            float(np.real(np.trace(rho@pz))))

def plot_histogram(circuit, shots=1024):
    backend = StatevectorSimulator()
    if circuit.num_qubits > 10: shots = 100
    result = execute(circuit, backend, shots=shots)
    counts = result.get_counts()
    states = [f"|{i:0{circuit.num_qubits}b}>" for i in range(min(2**circuit.num_qubits,64))]
    probs  = [counts.get(f"{i:0{circuit.num_qubits}b}",0)/shots for i in range(min(2**circuit.num_qubits,64))]
    fig, ax = plt.subplots(figsize=(10,5))
    ax.bar(states, probs, color="steelblue")
    ax.set_xlabel("State"); ax.set_ylabel("Probability")
    ax.set_title("PKTron Quantum Histogram", color="darkblue")
    plt.xticks(rotation=45, ha="right"); plt.tight_layout(); plt.show()
    return counts

def plot_histogram_interactive(circuits, shots=1024, labels=None):
    backend = StatevectorSimulator()
    if not isinstance(circuits, list): circuits = [circuits]
    if labels is None: labels = [f"Circuit {i+1}" for i in range(len(circuits))]
    fig = go.Figure()
    for circuit, label in zip(circuits, labels):
        if circuit.num_qubits > 10: shots = 100
        result = execute(circuit, backend, shots=shots)
        counts = result.get_counts()
        states = sorted(counts.keys())
        probs = [counts[s]/shots for s in states]
        fig.add_trace(go.Bar(x=states, y=probs, name=label))
    fig.update_layout(title="PKTron Interactive Histogram", xaxis_title="State", yaxis_title="Probability")
    return fig

def plot_bloch(circuit):
    if circuit.backend_type != "statevector":
        raise ValueError("Bloch requires statevector backend")
    psi = circuit.get_state()
    rho = np.outer(psi, psi.conj())
    n = circuit.num_qubits
    cols = min(n,3); rows = (n//cols)+(n%cols>0)
    fig = plt.figure(figsize=(6*cols, 6*rows))
    for q in range(n):
        try:
            rho_q = partial_trace(rho, [q])
            xv,yv,zv = get_bloch_vector(rho_q)
            ax = fig.add_subplot(rows, cols, q+1, projection="3d")
            u = np.linspace(0,2*np.pi,20); v = np.linspace(0,np.pi,20)
            ax.plot_surface(np.outer(np.cos(u),np.sin(v)),
                            np.outer(np.sin(u),np.sin(v)),
                            np.outer(np.ones(20),np.cos(v)), alpha=0.1, color="b")
            ax.scatter([xv],[yv],[zv], color="red", s=100)
            ax.plot([0,xv],[0,yv],[0,zv], color="red", linewidth=2)
            ax.set_title(f"Qubit {q}")
        except Exception as e:
            print(f"Bloch warning qubit {q}: {e}")
    fig.suptitle("PKTron Bloch Spheres", fontsize=14, color="darkblue")
    plt.tight_layout(); plt.show()
    return fig

def plot_bloch_interactive(circuit, animate=False):
    if circuit.backend_type != "statevector":
        raise ValueError("Bloch requires statevector backend")
    psi = circuit.get_state()
    rho = np.outer(psi, psi.conj())
    n = circuit.num_qubits
    if n > 5: raise ValueError("Interactive Bloch supports up to 5 qubits")
    cols = min(n,3); rows = (n//cols)+(n%cols>0)
    fig = make_subplots(rows=rows, cols=cols,
                        specs=[[{"type":"scene"}]*cols]*rows,
                        subplot_titles=[f"Qubit {i}" for i in range(n)])
    th = np.linspace(0,np.pi,30); ph = np.linspace(0,2*np.pi,30)
    th,ph = np.meshgrid(th,ph)
    xs=np.sin(th)*np.cos(ph); ys=np.sin(th)*np.sin(ph); zs=np.cos(th)
    for q in range(n):
        rho_q = partial_trace(rho,[q])
        xv,yv,zv = get_bloch_vector(rho_q)
        r,c = (q//cols)+1,(q%cols)+1
        fig.add_trace(go.Surface(x=xs,y=ys,z=zs,opacity=0.2,showscale=False),row=r,col=c)
        fig.add_trace(go.Scatter3d(x=[0,xv],y=[0,yv],z=[0,zv],mode="lines",
                                   line=dict(color="red",width=5)),row=r,col=c)
        fig.add_trace(go.Scatter3d(x=[xv],y=[yv],z=[zv],mode="markers",
                                   marker=dict(size=5,color="red")),row=r,col=c)
    fig.update_layout(title="PKTron Interactive Bloch", height=300*rows, width=300*cols)
    return fig

def plot_error_analysis(circuit, shots=1024, noise_prob=0.1):
    counts_list = []
    for _ in range(10):
        noisy = {}
        for _ in range(shots):
            tmp = QuantumCircuit(circuit.num_qubits,"statevector")
            for gate in circuit.gate_history:
                parts = gate.split()
                if parts[0]=="H": tmp.h(int(parts[-1]))
                elif parts[0]=="X": tmp.x(int(parts[-1]))
                elif parts[0]=="CNOT": tmp.cnot(int(parts[2]),int(parts[4]))
                if np.random.random()<noise_prob:
                    tmp.x(np.random.randint(0,circuit.num_qubits))
            o,_ = tmp.measure()
            k = format(o,f"0{circuit.num_qubits}b")
            noisy[k] = noisy.get(k,0)+1
        counts_list.append(noisy)
    states = sorted(set().union(*[set(c.keys()) for c in counts_list]))
    data = [[c.get(s,0)/shots for c in counts_list] for s in states]
    plt.figure(figsize=(10,6))
    sns.violinplot(data=data)
    plt.xticks(range(len(states)),states,rotation=45)
    plt.title("PKTron Error Analysis"); plt.tight_layout(); plt.show()
    return data

def plot_density_matrix(circuit, save_gif=False):
    if circuit.backend_type != "statevector":
        raise ValueError("Density matrix requires statevector backend")
    psi = circuit.get_state()
    rho = np.outer(psi, psi.conj())
    fig, ax = plt.subplots(figsize=(8,6))
    sns.heatmap(np.abs(rho), cmap="viridis", ax=ax)
    ax.set_title("PKTron Density Matrix")
    plt.show()
    return fig
