import os

import kp
import numpy as np
import logging
import pyshader as ps

DIRNAME = os.path.dirname(os.path.abspath(__file__))

def test_opalgobase_file():
    """
    Test basic OpMult operation
    """

    tensor_in_a = kp.Tensor([2, 2, 2])
    tensor_in_b = kp.Tensor([1, 2, 3])
    tensor_out = kp.Tensor([0, 0, 0])

    mgr = kp.Manager()
    mgr.rebuild([tensor_in_a, tensor_in_b, tensor_out])

    shader_path = os.path.join(DIRNAME, "../../shaders/glsl/opmult.comp.spv")

    mgr.eval_algo_file_def([tensor_in_a, tensor_in_b, tensor_out], shader_path)

    mgr.eval_tensor_sync_local_def([tensor_out])

    assert tensor_out.data() == [2.0, 4.0, 6.0]


params = [kp.Tensor([2, 2, 2]), kp.Tensor([1, 2, 3]), kp.Tensor([0, 0, 0])]

mgr = kp.Manager()
op_ct = kp.OpTensorCreate(params)
op_ct = mgr.rebuild(op_ct)
mgr.eval_op(op_ct)

algo = kp.Algo(params, spirv)
op_ac = kp.OpAlgoCreate(algo)
op_ac = mgr.rebuild(op_ac)
mgr.eval_op(op_ac)

op_ac = kp.OpAlgoCreate(kp.Algo(params, spirv))
mgr.eval_op(kp.OpAlgoCreate(algo))


mgr = kp.Manager()

op_ct = kp.OpTensorCreate(mgr, params) # This initialises operation
op_ct.eval()

algo = kp.Algo(params, spirv)
op_ac = kp.OpAlgoCreate(mgr, algo)
op_ct.eval()

op_tsd = kp.OpTensorSyncDevice(mgr, params)
op_ad = kp.OpAlgoDispatch(mgr, algo)
op_tsl = kp.OpTensorSyncLocal(mgr, params)

sq = kp.Sequence(mgr, "newSeq")
sq.record([op_tsd, op_ad, op_tsl])
sq.eval()
sq.destroy()

# Explore consistent interface:
op_tsd = kp.OpTensorSyncDevice(sq, params)
op_ad = kp.OpAlgoDispatch(sq, algo)
op_tsl = kp.OpTensorSyncLocal(sq, params)

op_tsd.record()
op_ad.record()
op_tsl.record()

sq.eval()



auto params = ...;
std::string shader = "...";
std::vector<uint32_t> spirv = kp::Shader::compile_source(shader);

// Example passing mgr
kp::Manager mgr;

kp::OpTensorCreate op_tc(mgr, params);
op_tc.eval()

kp::Algorithm algo(params, spirv);
kp::OpAlgoCreate op_ac(mgr, algo);
op_ac.eval()

op_ac.destroy()
op_tc.destroy()

kp::OpTensorAlgoCreate op_c(mgr, params, algo);
op_c.eval()

kp::Sequence sq(mgr);

kp::OpTensorSyncDevice op_tsd(mgr, params);
kp::OpAlgoDispatch op_ad(mgr, algo);
kp::OpTensorSyncLocal op_tsl(mgr, params);

sq.record({op_tsd, op_ad, op_tsl})

for(...) {
    sq.eval();

    tensorA...
}

######
#######
#######
#######
#######
######
// Example not passing mgr
kp::Manager mgr;

std::shared_ptr<kp::OpTensorCreate> op_tc_1{ new kp::OpTensorCreate(params) };
auto sq_1 = mgr.eval(op_tc_1); // Initialises and stores op as part of new sequence
mgr.eval(op_tc_1); // Fails as this op can only be "initialised" once
mgr.destroy(op_tc_1);
mgr.eval(op_tc_1); // This works as it's a new setup
mgr.eval<kp::OpTensorCreate>(params); // Fails as tensors already created
// NOT ALLOED TO DELETE JUST TENSORS ANYMORE - SEE BELOW
mgr.destroy(params); // Sends to inconsistent state as op_tc_1 will still destroy these parameters
mgr.destroy(op_tc_1, recursive=false); // Destroys only operation, which is useful when you need to ensure another operation owns the parameters
auto op_tc_2 = mgr.eval<kp::OpTensorCreate>(params);
std::shared_ptr<kp::OpTensorCreate> op_tc_2{ new kp::OpTensorCreate(params) }; // fails as tensors already created
op_tc_2.destroy(); // Manager still holds dangling reference so requires explicit termination in manager
mgr.destroy(op_tc_2);
auto op_tc_3 = mgr.eval({ new kp::OpTensorCreate(params) });

std::shared_ptr<kp::Algorithm> algo{ new kp::Algorithm(params, spirv, kp::Workgroup(), kp::SpecConst(), kp::PushConst()) };
std::shared_ptr<kp::OpAlgoCreate> op_ac_1{ new kp::OpAlgoCreate(algo) };
mgr.eval(op_ac_1); // Initialises and stores op as part of manager
mgr.eval(op_ac_1); // Fails as this op can only be "initialised" once
mgr.destroy(op_ac_1);

std::shared_ptr<kp::OpAlgoCreate> op_ac_2 =
    mgr.eval({ new kp::OpAlgoCreate(params, { new kp::Algorithm(spirv) }) });

std::shared_ptr<kp::OpAlgoMultCreate> op_amc{ new kp::OpAlgoMultCreate(params) };
mgr.eval(op_amc);

std::shared_ptr<kp::Algorithm> algo_mult = op_amc.algorithm()
std::vector<std::shared_ptr<kp::Tensor>> params = op_amc.tensors()

auto op_tsd = std::make_shared<kp::OpTensorSyncDevice>(params);
auto op_ad = std::make_shared<kp::OpAlgoSetPushConst>(algo);
auto op_ad = std::make_shared<kp::OpAlgoDispatch>(algo);
auto op_tsl = std::make_shared<kp::OpTensorSyncLocal>(mgr, params);

op_params = {op_tsd, op_ad, op_tsl};

mgr.record(op_params);
mgr.eval(); // Runs recorded default sequence

mgr.record(op_params, clear=false); // Non-create ops ok if rerun
mgr.eval(); // Runs twice the recorded paams

mgr.record("namedSeq", op_params);
mgr.eval("namedSeq");

kp::Manager mgrAsync(0, {0, 2});
mgr.sequence("namedSeq2", 0); // Create named sequence with queue in index 0
mgr.sequence("namedSeq3", 1);

mgr.eval_async("namedSeq2", op_params); // Clear, record params and eval
mgr.eval_async("namedSeq3", op_params); // Clear, record params and eval

mgr.eval_await("namedSeq2");
mgr.eval_await("namedSeq3");

mgr.destroy("namedSeq"); // Destroy named sequence
mgr.destroy({"namedSeq2", "namedSeq3"}); // Destroy multiple named sequences
mgr.destroy("namedSeq"); // Error




mgr = kp.Manager(0, [0, 2])

// Manager does not need to manage seq anymore
sq_1 = kp.Sequence(mgr, 0)

t1 = kp.Tensor(sq_1, [0, 0, 0])
t2 = kp.Tensor(sq_1, [0, 1, 2])

algo = kp.Algorithm(sq_1)

op_tc = kp.OpTensorCreate(sq_1, params)
op_tsd = kp.OpTensorSyncDevice(sq_1, params)
op_ac = kp.OpAlgoCreate(sq_1, algo)
op_ad = kp.OpAlgoDispatch(sq_1, algo)

sq_1.clear()

op_tc.record()
op_tsd.record()
op_ac.record()
op_ad.record()
op_ad.record()
op_ad.record()

sq_1.eval()


std::shared_ptr<kp::Manager> mgr = kp::ManagerSP(0, {0, 1});

std::shared_ptr<kp::Sequence> sq_2 = kp::SequenceSP(mgr, 1)

std::shared_ptr<kp::Tensor> t1 = kp::TensorSP(sq_2, {1, 2, 3});
std::shared_ptr<kp::Tensor> t2 = kp::TensorSP(sq_2, {2, 3, 4});

auto params = ...

std::shared_ptr<kp::Algorithm> algo2 = kp::AlgorithmSP(sq_2, params, spirv, workgroup);

// How do we deal with this?
{
    auto op_1 = kp::OpTensorSyncDevice(sq_2, params)
    auto op_2 = kp::OpAlgoDispatch(sq_2, algo)
}

sq_2.eval()


// HEAP ONLY - This would fail

kp::Manager mgr = kp::Manager(0, {0, 1});

kp::Sequence sq_2 = kp::Sequence(mgr, 1)

kp::Tensor t1 = kp::Tensor(sq_2, {1, 2, 3});
kp::Tensor t2 = kp::Tensor(sq_2, {2, 3, 4});

auto params = ...

kp::Algorithm algo2 = kp::AlgorithmSP(sq_2, params, spirv, workgroup);

// How do we deal with this?
{
    auto op_1 = kp::OpTensorSyncDevice(sq_2, params)
    auto op_2 = kp::OpAlgoDispatch(sq_2, algo)
}

sq_2.eval()





kp::Manager mgr = kp::Manager(0, {0, 1});

kp::Sequence sq_2 = kp::Sequence(mgr, 1)

kp::Tensor t1 = kp::Tensor(sq_2, {1, 2, 3});
kp::Tensor t2 = kp::Tensor(sq_2, {2, 3, 4});

auto params = ...

kp::Algorithm* algo2 = new kp::Algorithm(sq_2, params, spirv, workgroup);

// How do we deal with this?
{
    auto op_1 = kp::OpTensorSyncDevice(sq_2, params)
    auto op_2 = kp::OpAlgoDispatch(sq_2, algo)
}

sq_2.eval()






kp::Manager mgr = kp::Manager;

auto sq_2 = mgr.sequence()

{
    // What if we want to use tensor in a different sequence?
    auto t1 = sq_2.tensor({1, 2, 3});
    auto t2 = sq_2.tensor({1, 2, 3});

    auto algo2 = sq_2.algorithm();

    sq_2.record(kp::OpTensorRebuild({ t1 }))
    sq_2.record(kp::OpAlgoRebuild(params, algo2, spirv))
    sq_2.record(kp::OpTensorSyncDevice(prams))
    sq_2.record(kp::OpAlgoDispatch(prams, algo2))
}

sq_2.eval()



kp::Manager mgr = kp::Manager;

auto t1 = mgr.tensor({1, 2, 3}); // Held as weak ptr but passed as shared
auto t2 = mgr.tensor({1, 2, 3});

auto algo2 = mgr.algorithm();

{
    auto sq_2 = mgr.sequence()

    {
        sq_2.record(kp::OpTensorRebuild({ t1 })) // record only supports move operator &&
        sq_2.record(kp::OpAlgoRebuild(params, algo2, spirv))
        sq_2.record(kp::OpTensorSyncDevice(prams))
        sq_2.record(kp::OpAlgoDispatch(prams, algo2))
    }

    sq_2.eval()
}



// What about only tensors being init with it


{
    kp::Manager mgr = kp::Manager;

    auto t0 = mgr.tensor({0, 0, 0})

    {
        auto t1 = mgr.tensor({1, 2, 3}); // Held as weak ptr but passed as shared (refc 1)

        {
            auto sq_2 = mgr.sequence()

            {

                auto t2 = mgr.tensor({1, 2, 3}); // Held as weak ptr but passed as shared (refc 1)
                auto algo2 = mgr.algorithm(); // Held as weak ptr but passed as shared (refc 1)

                params = {t1, t2}

                sq_2.record(kp::OpTensorRebuild(params, {1, 2, 3, 4})) // Refc is now 2 for 3 for params
                sq_2.record(kp::OpAlgoRebuild(params, algo2, spirv)) // refc is now 2 for algo2, 3 for parms
                sq_2.record(kp::OpTensorSyncDevice(prams)) // refc for params 4
                sq_2.record(kp::OpAlgoDispatch(prams, algo2)) // refc for params 5, 3 for algo2
            }

            sq_2.eval() // all refcs stil valid
        } // seq destroyed so refc for algo2 and t2 drops to 0, gets destroyed, t1 has 1
    } // t1 refc drops to 0, gets destroyed
    // refc of t0 is still 1

    mgr.gc() // Iterates through all tensor, sequence and algo weak_ptr and removes unused

    // can we have something like
    mgr.sequence()
        .record(kp::OpTensorRebuild(params, {1, 2, 3, 4}))
        .record(kp::OpAlgoDispatch(params, algo2))
        .eval();

}// refc is destroyed by manager manually, the rest are empty shells so ignored




kp::Manager mgr = kp::Manager(0, {0, 1});

std::shared_ptr<kp::Tensor> t1 = mgr.tensor({1, 2, 3});
std::shared_ptr<kp::Tensor> t2 = mgr.tensor({1, 2, 3});

auto params = ...

std::shared_ptr<kp::Algorithm> algo2 = mgr.algorithm(params, spirv, workgroup);

sq_2.record<kp::OpTensorSyncDevice>(prams)
sq_2.record<kp::OpAlgoDispatch>(algo)


// WHY NO MORE DETROY TENSORS:

     * std::shared_ptr<kp::OpTensorCreate> op_tc1{ kp::OpTensorCreate(params) };
     * {
     *     std::shared_ptr<kp::OpTensorCreate> op_tc2{ kp::OpTensorCreate(params) };
     *     mgr.eval(op_tc2);
     *     mgr.destroy(params);
     *
     *     mgr.eval(op_tc1);
     *
     * } // op_tc1 is destroyed and all parameters are freed



// NO LONGER ALLOWED:  Mainly as manager now needs to regsiter ops
// If we still want it, then sequence wil have to hold ref to manager
auto sq = mgr.sequence();

auto op_tsd = std::make_shared<kp::OpTensorSyncDevice>(params);
auto op_ad = std::make_shared<kp::OpAlgoDispatch>(algo);
auto op_tsl = std::make_shared<kp::OpTensorSyncLocal>(mgr, params);

sq.record({op_tsd, op_ad, op_tsl}); // Clear and record
sq.eval();
sq.record({op_tsd, op_ad, op_tsl}, clear=false); // record on top
sq.eval();
sq.clear(); // explicitly clear





mgr = kp.Manager()

op_ct = kp.OpTensorCreate(params)
mgr.eval(op_ct)

algo = kp.Algo(params, spirv)
op_ac = kp.OpAlgoCreate(algo)
mgr.eval(op_ac) # Runs init on operator function (below shows explicit steps)

op_tsd = kp.OpTensorSyncDevice(params)
op_ad = kp.OpAlgoDispatch(algo)
op_tsl = kp.OpTensorSyncLocal(params)

sq = mgr.sequence()
sq.record([op_tsd, op_ad, op_tsl])
sq.eval()
sq.eval()
sq.eval()

mgr.eval(op_ac) # Would fail as algo is initialised
mgr.destroy(op_ac) # Destroys Op and Algo owned object
mgr.eval(op_ac) # Succeeds with new
mgr.destroy(op_ac)
mgr.init(op_ac)
mgr.eval(op_ac, init=False)











def test_shader_str():
    """
    Test basic OpAlgoBase operation
    """

    shader = """
#version 450
layout(set = 0, binding = 0) buffer tensorLhs {float valuesLhs[];};
layout(set = 0, binding = 1) buffer tensorRhs {float valuesRhs[];};
layout(set = 0, binding = 2) buffer tensorOutput { float valuesOutput[];};
layout (local_size_x = 1, local_size_y = 1, local_size_z = 1) in;

void main()
{
    uint index = gl_GlobalInvocationID.x;
    valuesOutput[index] = valuesLhs[index] * valuesRhs[index];
}
    """

    tensor_in_a = kp.Tensor([2, 2, 2])
    tensor_in_b = kp.Tensor([1, 2, 3])
    tensor_out = kp.Tensor([0, 0, 0])

    mgr = kp.Manager()
    mgr.rebuild([tensor_in_a, tensor_in_b, tensor_out])

    spirv = kp.Shader.compile_source(shader)

    mgr.eval_algo_data_def([tensor_in_a, tensor_in_b, tensor_out], spirv)

    mgr.eval_tensor_sync_local_def([tensor_out])

    assert tensor_out.data() == [2.0, 4.0, 6.0]

def test_sequence():
    """
    Test basic OpAlgoBase operation
    """
    mgr = kp.Manager(0, [2])

    tensor_in_a = kp.Tensor([2, 2, 2])
    tensor_in_b = kp.Tensor([1, 2, 3])
    tensor_out = kp.Tensor([0, 0, 0])

    mgr.rebuild([tensor_in_a, tensor_in_b, tensor_out])

    shader_path = os.path.abspath(os.path.join(DIRNAME, "../../shaders/glsl/opmult.comp.spv"))
    mgr.eval_async_algo_file_def([tensor_in_a, tensor_in_b, tensor_out], shader_path)

    mgr.eval_await_def()

    seq = mgr.sequence("op")
    seq.begin()
    seq.record_tensor_sync_local([tensor_in_a])
    seq.record_tensor_sync_local([tensor_in_b])
    seq.record_tensor_sync_local([tensor_out])
    seq.end()
    seq.eval()

    mgr.destroy("op")

    assert seq.is_init() == False

    assert tensor_out.data() == [2.0, 4.0, 6.0]
    assert np.all(tensor_out.numpy() == [2.0, 4.0, 6.0])

    mgr.destroy(tensor_in_a)
    mgr.destroy([tensor_in_b, tensor_out])

    assert tensor_in_a.is_init() == False
    assert tensor_in_b.is_init() == False
    assert tensor_out.is_init() == False

def test_workgroup():
    mgr = kp.Manager(0)

    tensor_a = kp.Tensor(np.zeros([16,8]))
    tensor_b = kp.Tensor(np.zeros([16,8]))

    mgr.rebuild([tensor_a, tensor_b])

    @ps.python2shader
    def compute_shader_wg(gl_idx=("input", "GlobalInvocationId", ps.ivec3),
                          gl_wg_id=("input", "WorkgroupId", ps.ivec3),
                          gl_wg_num=("input", "NumWorkgroups", ps.ivec3),
                          data1=("buffer", 0, ps.Array(ps.f32)),
                          data2=("buffer", 1, ps.Array(ps.f32))):
        i = gl_wg_id.x * gl_wg_num.y + gl_wg_id.y
        data1[i] = f32(gl_idx.x)
        data2[i] = f32(gl_idx.y)

    seq = mgr.sequence("new")
    seq.begin()
    seq.record_algo_data([tensor_a, tensor_b], compute_shader_wg.to_spirv(), workgroup=(16,8,1))
    seq.end()
    seq.eval()

    mgr.destroy(seq)

    assert seq.is_init() == False

    mgr.eval_tensor_sync_local_def([tensor_a, tensor_b])

    print(tensor_a.numpy())
    print(tensor_b.numpy())

    assert np.all(tensor_a.numpy() == np.stack([np.arange(16)]*8, axis=1).ravel())
    assert np.all(tensor_b.numpy() == np.stack([np.arange(8)]*16, axis=0).ravel())

    mgr.destroy([tensor_a, tensor_b])

    assert tensor_a.is_init() == False
    assert tensor_b.is_init() == False


def test_tensor_rebuild_backwards_compat():
    """
    Test basic OpMult operation
    """

    tensor_in_a = kp.Tensor([2, 2, 2])
    tensor_in_b = kp.Tensor([1, 2, 3])
    tensor_out = kp.Tensor([0, 0, 0])

    mgr = kp.Manager()

    mgr.eval_tensor_create_def([tensor_in_a, tensor_in_b, tensor_out])

    shader_path = os.path.abspath(os.path.join(DIRNAME, "../../shaders/glsl/opmult.comp.spv"))
    mgr.eval_async_algo_file_def([tensor_in_a, tensor_in_b, tensor_out], shader_path)
    mgr.eval_await_def()

    mgr.eval_tensor_sync_local_def([tensor_out])

    assert tensor_out.data() == [2.0, 4.0, 6.0]
    assert np.all(tensor_out.numpy() == [2.0, 4.0, 6.0])