/*//////////////////////////////////////////////////////////////////////////////
// CMT Cosmic Muon Tomography project //////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////

  Copyright (c) 2014, Universita' degli Studi di Padova, INFN sez. di Padova
  All rights reserved

  Authors: Andrea Rigoni Garola < andrea.rigoni@pd.infn.it >

  ------------------------------------------------------------------
  This library is free software;  you  can  redistribute  it  and/or
  modify it  under the  terms  of  the  GNU  Lesser  General  Public
  License as published  by  the  Free  Software  Foundation;  either
  version 3.0 of the License, or (at your option) any later version.

  This library is  distributed in  the hope that it will  be useful,
  but  WITHOUT ANY WARRANTY;  without  even  the implied warranty of
  MERCHANTABILITY  or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  Lesser General Public License for more details.

  You should have received a copy of  the GNU Lesser General  Public
  License along with this library.

//////////////////////////////////////////////////////////////////////////////*/

#include "Threads.h"
#include <chrono>

#ifdef _OPENMP
#include <omp.h>
#endif

#ifdef __linux__
#include <pthread.h>
#include <sched.h>
#endif

namespace uLib {

Thread::Thread() : m_Running(false) {}

Thread::~Thread() {
  if (m_Thread.joinable()) {
    m_Thread.detach();
  }
}

void Thread::Start() {
  Mutex::ScopedLock lock(m_ThreadMutex);
  if (m_Running) return;
  
  m_Running = true;
  m_Thread = std::thread(&Thread::ThreadEntryPoint, this);
}

void Thread::Join() {
  if (m_Thread.joinable()) {
    m_Thread.join();
  }
}

void Thread::Detach() {
  if (m_Thread.joinable()) {
    m_Thread.detach();
  }
}

bool Thread::IsJoinable() const {
  return m_Thread.joinable();
}

bool Thread::IsRunning() const {
  return m_Running;
}

void Thread::Run() {
  // Override in subclasses
}

void Thread::ThreadEntryPoint() {
  this->Run();
  m_Running = false;
}

void Thread::Sleep(int milliseconds) {
  std::this_thread::sleep_for(std::chrono::milliseconds(milliseconds));
}

void Thread::Yield() {
  std::this_thread::yield();
}

void Thread::SetAffinity(int cpu) {
#ifdef __linux__
  if (m_Thread.joinable()) {
    cpu_set_t cpuset;
    CPU_ZERO(&cpuset);
    CPU_SET(cpu, &cpuset);
    pthread_setaffinity_np(m_Thread.native_handle(), sizeof(cpu_set_t), &cpuset);
  }
#endif
}

void Thread::SetAffinity(const std::vector<int>& cpus) {
#ifdef __linux__
  if (m_Thread.joinable()) {
    cpu_set_t cpuset;
    CPU_ZERO(&cpuset);
    for (int cpu : cpus) {
        CPU_SET(cpu, &cpuset);
    }
    pthread_setaffinity_np(m_Thread.native_handle(), sizeof(cpu_set_t), &cpuset);
  }
#endif
}

void Thread::SetNumThreads(int n) {
#ifdef _OPENMP
  omp_set_num_threads(n);
#endif
}

int Thread::GetNumThreads() {
#ifdef _OPENMP
  return omp_get_max_threads();
#else
  return 1;
#endif
}

int Thread::GetThreadNum() {
#ifdef _OPENMP
  return omp_get_thread_num();
#else
  return 0;
#endif
}

// Team Implementation //

Team::Team(int num_threads) : m_Size(num_threads), m_UseOpenMP(false) {
#ifdef _OPENMP
  m_UseOpenMP = true;
  if (m_Size > 0) omp_set_num_threads(m_Size);
  else m_Size = omp_get_max_threads();
#else
  if (m_Size <= 0) m_Size = 1;
#endif
}

Team::~Team() {
  Wait();
}

void Team::Run(Task* task) {
  if (!task) return;
#ifdef _OPENMP
  if (m_UseOpenMP) {
    #pragma omp task
    task->Execute();
    return;
  }
#endif
  // Fallback to synchronous execution if no OpenMP
  task->Execute();
}

void Team::Wait() {
#ifdef _OPENMP
  if (m_UseOpenMP) {
    #pragma omp taskwait
  }
#endif
}

void Team::SetSize(int n) {
  m_Size = n;
#ifdef _OPENMP
  if (m_UseOpenMP) omp_set_num_threads(m_Size);
#endif
}

void Team::SetAffinity(const std::vector<int>& cpus) {
  if (cpus.empty()) return;
#ifdef __linux__
#ifdef _OPENMP
  if (m_UseOpenMP) {
    #pragma omp parallel
    {
      int tid = omp_get_thread_num();
      int cpu = cpus[tid % cpus.size()];
      cpu_set_t cpuset;
      CPU_ZERO(&cpuset);
      CPU_SET(cpu, &cpuset);
      pthread_setaffinity_np(pthread_self(), sizeof(cpu_set_t), &cpuset);
    }
  }
#endif
#endif
}

} // namespace uLib